From 39d9ffe1db2217efd2377b6a6d7e42ac5495def0 Mon Sep 17 00:00:00 2001 From: William Desportes Date: Tue, 24 Sep 2024 00:15:34 +0200 Subject: [PATCH] Make a working client server --- snow-scanner/src/server.rs | 35 +++++++++---- snow-scanner/src/worker/modules.rs | 80 +++++++++++++++++++++++++++--- snow-scanner/src/worker/worker.rs | 57 +++++++++++++++++---- 3 files changed, 143 insertions(+), 29 deletions(-) diff --git a/snow-scanner/src/server.rs b/snow-scanner/src/server.rs index 1541c26..5678095 100644 --- a/snow-scanner/src/server.rs +++ b/snow-scanner/src/server.rs @@ -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, @@ -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 = 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 } } diff --git a/snow-scanner/src/worker/modules.rs b/snow-scanner/src/worker/modules.rs index ee0a97e..9ae42be 100644 --- a/snow-scanner/src/worker/modules.rs +++ b/snow-scanner/src/worker/modules.rs @@ -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 }, #[serde(rename = "")] - Invalid, + Invalid { err: String }, +} + +impl<'de> Deserialize<'de> for Network { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = ::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(&self, serializer: S) -> Result + 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 for String { fn into(self) -> WorkerMessages { - let req: Result = - serde_json::from_str(self.as_str()); + let req: Result = 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() + } + ); + } +} diff --git a/snow-scanner/src/worker/worker.rs b/snow-scanner/src/worker/worker.rs index 40d9926..43aed0d 100644 --- a/snow-scanner/src/worker/worker.rs +++ b/snow-scanner/src/worker/worker.rs @@ -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, + pub last_request_for_work: Option, } impl Worker { @@ -17,6 +25,8 @@ impl Worker { info!("New worker"); Worker { authenticated: false, + tasks: vec![], + last_request_for_work: None, } } @@ -48,11 +58,20 @@ 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?)*/ - Ok(()) + 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(()) + } + } } } @@ -81,19 +100,35 @@ fn main() -> () { Ok(_) => {} Err(err) => error!("Processing error: {err}"), } + let mut request: Option = 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), } }