wdes/security
1
0
Fork 0
Code Actions Packages Releases Activity

Make a working client server

Browse Source
This commit is contained in:
William Desportes
2024-09-24 00:15:34 +02:00
parent 27c3f7ecd1
commit 39d9ffe1db
3 changed files with 143 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
snow-scanner/src/server.rs
View File

@ -1,8 +1,9 @@
use cidr::IpCidr;
use log2::*;
use std::collections::HashMap;
use std::{collections::HashMap, str::FromStr};
use ws2::{Pod, WebSocket};
use crate::worker::modules::WorkerMessages;
use crate::worker::modules::{Network, WorkerMessages};
pub struct Server {
pub clients: HashMap<u32, Worker>,
@ -76,15 +77,13 @@ impl ws2::Handler for Server {
info!("on message: {msg}, {ws}");
let worker_message: WorkerMessages = msg.clone().into();
let mut worker_reply: Option<WorkerMessages> = None;
let worker_request: WorkerMessages = msg.clone().into();
match worker_message {
let result = match worker_request {
WorkerMessages::AuthenticateRequest { login } => {
if !worker.is_authenticated() {
worker.authenticate(login);
/*let echo = format!("echo: {msg}");
let n = ws.send(echo);
return Ok(n?);*/
return Ok(());
} else {
error!("Already authenticated: {ws}");
@ -92,11 +91,12 @@ impl ws2::Handler for Server {
}
}
WorkerMessages::GetWorkRequest {} => {
let echo = format!("wr");
let n = ws.send(echo);
Ok(n?)
worker_reply = Some(WorkerMessages::DoWorkRequest {
neworks: vec![Network(IpCidr::from_str("127.0.0.0/31")?)],
});
Ok(())
}
WorkerMessages::Invalid => {
WorkerMessages::DoWorkRequest { .. } | WorkerMessages::Invalid { .. } => {
error!("Unable to understand: {msg}, {ws}");
// Unable to understand, close the connection
return ws.close();
@ -104,6 +104,19 @@ impl ws2::Handler for Server {
error!("No implemented: {:#?}", msg);
Ok(())
}*/
};
// it has a request to send
if let Some(worker_reply) = worker_reply {
let msg_string: String = worker_reply.to_string();
match ws.send(msg_string) {
Ok(_) => match worker_reply {
WorkerMessages::DoWorkRequest { .. } => {}
msg => error!("No implemented: {:#?}", msg),
},
Err(err) => error!("Error sending reply to {ws}: {err}"),
}
}
result
}
}

80
snow-scanner/src/worker/modules.rs
View File

@ -1,6 +1,12 @@
use serde::{Deserialize, Serialize};
use std::str::FromStr;
#[derive(Debug, Serialize, Deserialize, Clone)]
use cidr::IpCidr;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
#[derive(Debug, Clone, PartialEq)]
pub struct Network(pub IpCidr);
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(tag = "type", content = "request")]
pub enum WorkerMessages {
#[serde(rename = "auth_request")]
@ -8,9 +14,36 @@ pub enum WorkerMessages {
#[serde(rename = "get_work")]
GetWorkRequest {},
#[serde(rename = "do_work")]
DoWorkRequest {},
DoWorkRequest { neworks: Vec<Network> },
#[serde(rename = "")]
Invalid,
Invalid { err: String },
}
impl<'de> Deserialize<'de> for Network {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let s = <String>::deserialize(deserializer)?;
let k: &str = s.as_str();
match IpCidr::from_str(k) {
Ok(d) => Ok(Network(d)),
Err(err) => Err(serde::de::Error::custom(format!(
"Unsupported value {k}: {err}"
))),
}
}
}
impl Serialize for Network {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(self.0.to_string().as_str())
}
}
impl ToString for WorkerMessages {
@ -21,11 +54,44 @@ impl ToString for WorkerMessages {
impl Into<WorkerMessages> for String {
fn into(self) -> WorkerMessages {
let req: Result<WorkerMessages, serde_json::Error> =
serde_json::from_str(self.as_str());
let req: Result<WorkerMessages, serde_json::Error> = serde_json::from_str(self.as_str());
match req {
Ok(d) => d,
Err(_) => WorkerMessages::Invalid,
Err(err) => WorkerMessages::Invalid {
err: err.to_string(),
},
}
}
}
#[cfg(test)]
mod tests {
use cidr::IpCidr;
use super::*;
#[test]
fn deserialize_do_work_empty() {
let data = "{\"type\":\"do_work\",\"request\":{\"neworks\":[]}}";
let result: WorkerMessages = data.to_string().into();
assert_eq!(
result,
WorkerMessages::DoWorkRequest {
neworks: [].to_vec()
}
);
}
#[test]
fn deserialize_do_work() {
let data = "{\"type\":\"do_work\",\"request\":{\"neworks\":[\"127.0.0.0/31\"]}}";
let result: WorkerMessages = data.to_string().into();
let cidr: IpCidr = IpCidr::from_str("127.0.0.0/31").unwrap();
assert_eq!(
result,
WorkerMessages::DoWorkRequest {
neworks: [Network(cidr)].to_vec()
}
);
}
}

55
snow-scanner/src/worker/worker.rs
View File

@ -1,5 +1,6 @@
use std::env;
use std::{env, net::IpAddr};
use chrono::{Duration, NaiveDateTime, Utc};
use log2::*;
use ws2::{Pod, WebSocket};
@ -7,9 +8,16 @@ pub mod modules;
use crate::modules::WorkerMessages;
#[derive(Debug, Clone)]
pub struct IpToResolve {
pub address: IpAddr,
}
#[derive(Debug, Clone)]
pub struct Worker {
pub authenticated: bool,
pub tasks: Vec<IpToResolve>,
pub last_request_for_work: Option<NaiveDateTime>,
}
impl Worker {
@ -17,6 +25,8 @@ impl Worker {
info!("New worker");
Worker {
authenticated: false,
tasks: vec![],
last_request_for_work: None,
}
}
@ -48,12 +58,21 @@ impl ws2::Handler for Worker {
}
fn on_message(&mut self, ws: &WebSocket, msg: String) -> Pod {
/*info!("on message: {msg}, {ws}");
let echo = format!("echo: {msg}");
let n = ws.send(echo);
Ok(n?)*/
let worker_request: WorkerMessages = msg.clone().into();
match worker_request {
WorkerMessages::DoWorkRequest { neworks } => {
info!("Should work on: {:?}", neworks);
Ok(())
}
WorkerMessages::AuthenticateRequest { .. }
| WorkerMessages::GetWorkRequest {}
| WorkerMessages::Invalid { .. } => {
error!("Unable to understand message: {msg}, {:?}", worker_request);
Ok(())
}
}
}
}
fn main() -> () {
@ -81,19 +100,35 @@ fn main() -> () {
Ok(_) => {}
Err(err) => error!("Processing error: {err}"),
}
let mut request: Option<WorkerMessages> = None;
if !worker.is_authenticated() {
let msg: WorkerMessages = WorkerMessages::AuthenticateRequest {
request = Some(WorkerMessages::AuthenticateRequest {
login: "williamdes".to_string(),
};
let msg_string: String = msg.to_string();
});
} else {
if worker.last_request_for_work.is_none()
|| (worker.last_request_for_work.is_some()
&& Utc::now().naive_utc()
> (worker.last_request_for_work.unwrap() + Duration::minutes(10)))
{
request = Some(WorkerMessages::GetWorkRequest {});
}
}
// it has a request to send
if let Some(request) = request {
let msg_string: String = request.to_string();
match ws_client.send(msg_string) {
Ok(_) => {
match msg {
match request {
WorkerMessages::AuthenticateRequest { login } => {
worker.authenticated = true; // Anyway, it will kick us if this is not success
info!("Logged in as: {login}")
}
WorkerMessages::GetWorkRequest {} => {}
WorkerMessages::GetWorkRequest {} => {
worker.last_request_for_work = Some(Utc::now().naive_utc());
info!("Asked for work: {:?}", worker.last_request_for_work)
}
msg => error!("No implemented: {:#?}", msg),
}
}