From 58d6ed043e983d0012e12820019f5b8f8bf945d5 Mon Sep 17 00:00:00 2001 From: William Desportes Date: Mon, 23 Sep 2024 17:20:50 +0200 Subject: [PATCH] First working version of client and server --- snow-scanner/Cargo.toml | 10 +- snow-scanner/src/main.rs | 150 +++++++++------------------ snow-scanner/src/server.rs | 113 +++++++++++++++++++++ snow-scanner/src/worker/mod.rs | 1 + snow-scanner/src/worker/modules.rs | 9 ++ snow-scanner/src/worker/worker.rs | 158 +++++++++++++++++++++++++++++ 6 files changed, 334 insertions(+), 107 deletions(-) create mode 100644 snow-scanner/src/server.rs create mode 100644 snow-scanner/src/worker/mod.rs create mode 100644 snow-scanner/src/worker/modules.rs create mode 100644 snow-scanner/src/worker/worker.rs diff --git a/snow-scanner/Cargo.toml b/snow-scanner/Cargo.toml index 4ad8df7..612b073 100644 --- a/snow-scanner/Cargo.toml +++ b/snow-scanner/Cargo.toml @@ -29,14 +29,17 @@ maintenance = { status = "passively-maintained" } name = "snow-scanner" path = "src/main.rs" +[[bin]] +name = "snow-scanner-worker" +path = "src/worker/worker.rs" + # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +log2 = "0.1.11" +ws2 = "0.2.5" actix-web = "4" actix-files = "0.6.6" -hmac = "0.12.1" -sha2 = "0.10.8" -hex = "0.4.3" # mariadb-dev on Alpine # "mysqlclient-src" "mysql_backend" diesel = { version = "2.2.0", default-features = false, features = ["mysql", "chrono", "uuid", "r2d2"] } @@ -46,3 +49,4 @@ chrono = "0.4.38" uuid = { version = "1.10.0", default-features = false, features = ["v7", "serde", "std"] } cidr = "0.2.2" serde = "1.0.210" +serde_json = "1.0.128" diff --git a/snow-scanner/src/main.rs b/snow-scanner/src/main.rs index 72e2a7a..571d116 100644 --- a/snow-scanner/src/main.rs +++ b/snow-scanner/src/main.rs @@ -2,6 +2,7 @@ use actix_files::NamedFile; use actix_web::error::ErrorInternalServerError; use actix_web::http::header::ContentType; use actix_web::{web, App, HttpRequest, HttpResponse, HttpServer}; +use log2::*; use chrono::{NaiveDateTime, Utc}; use diesel::deserialize::{self, FromSqlRow}; @@ -11,6 +12,7 @@ use diesel::sql_types::Text; use diesel::r2d2::ConnectionManager; use diesel::r2d2::Pool; +use std::collections::HashMap; use std::io::Write; use std::path::PathBuf; use std::str::FromStr; @@ -30,8 +32,11 @@ use dns_ptr_resolver::{get_ptr, ResolvedResult}; pub mod models; pub mod schema; +pub mod server; +pub mod worker; use crate::models::*; +use crate::server::Server; /// Short-hand for the database pool type to use throughout the app. type DbPool = Pool>; @@ -262,8 +267,8 @@ async fn handle_scan(pool: web::Data, params: web::Form) -> ended_at: None, }; match scan_task.save(conn) { - Ok(_) => println!("Added {}", ip.to_string()), - Err(err) => eprintln!("Not added: {:?}", err), + Ok(_) => error!("Added {}", ip.to_string()), + Err(err) => error!("Not added: {:?}", err), } } }) @@ -533,16 +538,30 @@ async fn pong() -> HttpResponse { #[actix_web::main] async fn main() -> std::io::Result<()> { + let _log2 = log2::stdout() + .module(false) + .level(match env::var("RUST_LOG") { + Ok(level) => level, + Err(_) => "debug".to_string(), + }) + .start(); + let server_address: String = if let Ok(env) = env::var("SERVER_ADDRESS") { env } else { - "localhost:8000".to_string() + "127.0.0.1:8000".to_string() + }; + + let worker_server_address: String = if let Ok(env) = env::var("WORKER_SERVER_ADDRESS") { + env + } else { + "127.0.0.1:8800".to_string() }; let db_url: String = if let Ok(env) = env::var("DB_URL") { env } else { - eprintln!("Missing ENV: DB_URL"); + error!("Missing ENV: DB_URL"); "mysql://localhost".to_string() }; @@ -552,8 +571,8 @@ async fn main() -> std::io::Result<()> { let conn = &mut pool.get().unwrap(); let names = Scanner::list_names(Scanners::Stretchoid, conn); match names { - Ok(names) => println!("Found {} Stretchoid scanners", names.len()), - Err(err) => eprintln!("Unable to get names: {}", err), + Ok(names) => info!("Found {} Stretchoid scanners", names.len()), + Err(err) => error!("Unable to get names: {}", err), }; let server = HttpServer::new(move || { @@ -582,108 +601,31 @@ async fn main() -> std::io::Result<()> { .bind(&server_address); match server { Ok(server) => { - println!("Now listening on {}", server_address); + match ws2::listen(worker_server_address.as_str()) { + Ok(mut ws_server) => { + std::thread::spawn(move || { + let mut ws_server_handles = Server { + clients: HashMap::new(), + }; + info!("Worker server is listening on: {worker_server_address}"); + loop { + match ws_server.process(&mut ws_server_handles, 0.5) { + Ok(_) => {} + Err(err) => error!("Processing error: {err}"), + } + ws_server_handles.cleanup(&ws_server); + } + }); + } + Err(err) => error!("Unable to listen on {worker_server_address}: {err}"), + }; + + info!("Now listening on {}", server_address); server.run().await } Err(err) => { - eprintln!("Could not bind the server to {}", server_address); + error!("Could not bind the server to {}", server_address); Err(err) } } } -/* -(POST) (/register) => { - let data = try_or_400!(post_input!(request, { - email: String, - })); - - // We just print what was received on stdout. Of course in a real application - // you probably want to process the data, eg. store it in a database. - println!("Received data: {:?}", data); - - - let mut mac = HmacSha256::new_from_slice(b"my secret and secure key") - .expect("HMAC can take key of any size"); - mac.update(data.email.as_bytes()); - - // `result` has type `CtOutput` which is a thin wrapper around array of - // bytes for providing constant time equality check - let result = mac.finalize(); - // To get underlying array use `into_bytes`, but be careful, since - // incorrect use of the code value may permit timing attacks which defeats - // the security provided by the `CtOutput` - let code_bytes = result.into_bytes(); - rouille::Response::html(format!("Success! {}.", hex::encode(code_bytes))) -}, - -(GET) (/{api_key: String}/scanners/{scanner_name: String}) => { - let mut mac = HmacSha256::new_from_slice(b"my secret and secure key") - .expect("HMAC can take key of any size"); - - mac.update(b"williamdes@wdes.fr"); - - println!("{}", api_key); - let hex_key = hex::decode(&api_key).unwrap(); - // `verify_slice` will return `Ok(())` if code is correct, `Err(MacError)` otherwise - mac.verify_slice(&hex_key).unwrap(); - - rouille::Response::empty_404() -}, - - thread::spawn(move || { - let conn = &mut get_connection(db_url.as_str()); - // Reset scan tasks - let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, still_processing_at = NULL, started_at = NULL WHERE (still_processing_at IS NOT NULL OR started_at IS NOT NULL) AND ended_at IS NULL", - named_params! { - ":updated_at": Utc::now().naive_utc().to_string(), - }).unwrap(); - - loop { - let mut stmt = conn.prepare("SELECT task_group_id, cidr FROM scan_tasks WHERE started_at IS NULL ORDER BY created_at ASC").unwrap(); - let mut rows = stmt.query(named_params! {}).unwrap(); - println!("Waiting for jobs"); - while let Some(row) = rows.next().unwrap() { - let task_group_id: String = row.get(0).unwrap(); - let cidr_str: String = row.get(1).unwrap(); - let cidr: IpCidr = cidr_str.parse().expect("Should parse CIDR"); - println!("Picking up: {} -> {}", task_group_id, cidr); - println!("Range, from {} to {}", cidr.first(), cidr.last()); - let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, started_at = :started_at WHERE cidr = :cidr AND task_group_id = :task_group_id", - named_params! { - ":updated_at": Utc::now().naive_utc().to_string(), - ":started_at": Utc::now().naive_utc().to_string(), - ":cidr": cidr_str, - ":task_group_id": task_group_id, - }).unwrap(); - let addresses = cidr.iter().addresses(); - let count = addresses.count(); - let mut current = 0; - for addr in addresses { - match handle_ip(conn, addr.to_string()) { - Ok(scanner) => println!("Processed {}", scanner.ip), - Err(_) => println!("Processed {}", addr), - } - current += 1; - if (current / count) % 10 == 0 { - let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, still_processing_at = :still_processing_at WHERE cidr = :cidr AND task_group_id = :task_group_id", - named_params! { - ":updated_at": Utc::now().naive_utc().to_string(), - ":still_processing_at": Utc::now().naive_utc().to_string(), - ":cidr": cidr_str, - ":task_group_id": task_group_id, - }).unwrap(); - } - } - let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, ended_at = :ended_at WHERE cidr = :cidr AND task_group_id = :task_group_id", - named_params! { - ":updated_at": Utc::now().naive_utc().to_string(), - ":ended_at": Utc::now().naive_utc().to_string(), - ":cidr": cidr_str, - ":task_group_id": task_group_id, - }).unwrap(); - } - - let two_hundred_millis = Duration::from_millis(500); - thread::sleep(two_hundred_millis); - } - });*/ diff --git a/snow-scanner/src/server.rs b/snow-scanner/src/server.rs new file mode 100644 index 0000000..c140f6d --- /dev/null +++ b/snow-scanner/src/server.rs @@ -0,0 +1,113 @@ +use log2::*; +use std::collections::HashMap; +use ws2::{Pod, WebSocket}; + +use crate::worker::modules::WorkerMessages; + +pub struct Server { + pub clients: HashMap, +} + +impl Server { + pub fn cleanup(&self, _: &ws2::Server) -> &Server { + // TODO: implement check not logged in + &self + } +} + +#[derive(Debug, Clone)] +pub struct Worker { + pub authenticated: bool, + pub login: Option, +} + +impl Worker { + pub fn initial() -> Worker { + info!("New worker"); + Worker { + authenticated: false, + login: None, + } + } + + pub fn is_authenticated(&self) -> bool { + self.authenticated + } + + pub fn authenticate(&mut self, login: String) -> &Worker { + if self.authenticated { + warn!( + "Worker is already authenticated as {}", + self.login.clone().unwrap_or("".to_string()) + ); + return self; + } else { + info!("Worker is now authenticated as {login}"); + } + self.login = Some(login); + self.authenticated = true; + self + } +} + +impl ws2::Handler for Server { + fn on_open(&mut self, ws: &WebSocket) -> Pod { + info!("New client: {ws}"); + let worker = Worker::initial(); + // Add the client + self.clients.insert(ws.id(), worker); + Ok(()) + } + + fn on_close(&mut self, ws: &WebSocket) -> Pod { + info!("Client /quit: {ws}"); + // Drop the client + self.clients.remove(&ws.id()); + Ok(()) + } + + fn on_message(&mut self, ws: &WebSocket, msg: String) -> Pod { + let client = self.clients.get_mut(&ws.id()); + if client.is_none() { + // Impossible, close in case + return ws.close(); + } + let worker: &mut Worker = client.unwrap(); + + info!("on message: {msg}, {ws}"); + + let worker_message: WorkerMessages = match serde_json::from_str(msg.as_str()) { + Ok(worker_message) => worker_message, + Err(err) => { + error!("Unable to understand: {msg}, {ws}: {err}"); + // Unable to understand, close the connection + return ws.close(); + } + }; + + match worker_message { + 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: {msg}, {ws}"); + return Ok(()); + } + } + WorkerMessages::FooRequest { username } => { + let echo = format!("echo: {msg}"); + let n = ws.send(echo); + Ok(n?) + } + WorkerMessages::String => { + error!("Unable to understand: {msg}, {ws}"); + // Unable to understand, close the connection + return ws.close(); + } + } + } +} diff --git a/snow-scanner/src/worker/mod.rs b/snow-scanner/src/worker/mod.rs new file mode 100644 index 0000000..55e5f1d --- /dev/null +++ b/snow-scanner/src/worker/mod.rs @@ -0,0 +1 @@ +pub mod modules; diff --git a/snow-scanner/src/worker/modules.rs b/snow-scanner/src/worker/modules.rs new file mode 100644 index 0000000..e218dd6 --- /dev/null +++ b/snow-scanner/src/worker/modules.rs @@ -0,0 +1,9 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub enum WorkerMessages { + AuthenticateRequest { login: String }, + FooRequest { username: String }, + String, +} diff --git a/snow-scanner/src/worker/worker.rs b/snow-scanner/src/worker/worker.rs new file mode 100644 index 0000000..babc718 --- /dev/null +++ b/snow-scanner/src/worker/worker.rs @@ -0,0 +1,158 @@ +use std::env; + +use log2::*; +use ws2::{Pod, WebSocket}; + +pub mod modules; + +use crate::modules::WorkerMessages; + +#[derive(Debug, Clone)] +pub struct Worker { + pub authenticated: bool, +} + +impl Worker { + pub fn initial() -> Worker { + info!("New worker"); + Worker { + authenticated: false, + } + } + + pub fn is_authenticated(&self) -> bool { + self.authenticated + } + + pub fn authenticate(&mut self, login: String) -> &Worker { + if self.authenticated { + warn!("Worker is already authenticated"); + return self; + } else { + info!("Worker is now authenticated as {login}"); + } + self.authenticated = true; + self + } +} + +impl ws2::Handler for Worker { + fn on_open(&mut self, ws: &WebSocket) -> Pod { + info!("Connected to: {ws}, starting to work"); + Ok(()) + } + + fn on_close(&mut self, ws: &WebSocket) -> Pod { + info!("End of the work day: {ws}"); + Ok(()) + } + + 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(()) + } +} + +fn main() -> () { + let _log2 = log2::stdout() + .module(true) + .level(match env::var("RUST_LOG") { + Ok(level) => level, + Err(_) => "debug".to_string(), + }) + .start(); + info!("Running the worker"); + let url = "ws://127.0.0.1:8800"; + let mut worker = Worker::initial(); + match ws2::connect(url) { + Ok(mut ws_client) => { + let connected = ws_client.is_open(); + if connected { + info!("Connected to: {url}"); + } else { + info!("Connecting to: {url}"); + } + + loop { + match ws_client.process(&mut worker, 0.5) { + Ok(_) => {} + Err(err) => error!("Processing error: {err}"), + } + if ! worker.is_authenticated() { + let msg: WorkerMessages = WorkerMessages::AuthenticateRequest { + login: "williamdes".to_string(), + }; + let msg: String = serde_json::to_string(&msg).expect("To serialize").into(); + match ws_client.send(msg) { + Ok(_) => { + worker.authenticated = true; + } + Err(err) => error!("Unable to connect to {url}: {err}"), + } + } + } + } + Err(err) => error!("Unable to connect to {url}: {err}"), + } +} +/* +thread::spawn(move || { + let conn = &mut get_connection(db_url.as_str()); + // Reset scan tasks + let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, still_processing_at = NULL, started_at = NULL WHERE (still_processing_at IS NOT NULL OR started_at IS NOT NULL) AND ended_at IS NULL", + named_params! { + ":updated_at": Utc::now().naive_utc().to_string(), + }).unwrap(); + + loop { + let mut stmt = conn.prepare("SELECT task_group_id, cidr FROM scan_tasks WHERE started_at IS NULL ORDER BY created_at ASC").unwrap(); + let mut rows = stmt.query(named_params! {}).unwrap(); + println!("Waiting for jobs"); + while let Some(row) = rows.next().unwrap() { + let task_group_id: String = row.get(0).unwrap(); + let cidr_str: String = row.get(1).unwrap(); + let cidr: IpCidr = cidr_str.parse().expect("Should parse CIDR"); + println!("Picking up: {} -> {}", task_group_id, cidr); + println!("Range, from {} to {}", cidr.first(), cidr.last()); + let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, started_at = :started_at WHERE cidr = :cidr AND task_group_id = :task_group_id", + named_params! { + ":updated_at": Utc::now().naive_utc().to_string(), + ":started_at": Utc::now().naive_utc().to_string(), + ":cidr": cidr_str, + ":task_group_id": task_group_id, + }).unwrap(); + let addresses = cidr.iter().addresses(); + let count = addresses.count(); + let mut current = 0; + for addr in addresses { + match handle_ip(conn, addr.to_string()) { + Ok(scanner) => println!("Processed {}", scanner.ip), + Err(_) => println!("Processed {}", addr), + } + current += 1; + if (current / count) % 10 == 0 { + let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, still_processing_at = :still_processing_at WHERE cidr = :cidr AND task_group_id = :task_group_id", + named_params! { + ":updated_at": Utc::now().naive_utc().to_string(), + ":still_processing_at": Utc::now().naive_utc().to_string(), + ":cidr": cidr_str, + ":task_group_id": task_group_id, + }).unwrap(); + } + } + let _ = conn.execute("UPDATE scan_tasks SET updated_at = :updated_at, ended_at = :ended_at WHERE cidr = :cidr AND task_group_id = :task_group_id", + named_params! { + ":updated_at": Utc::now().naive_utc().to_string(), + ":ended_at": Utc::now().naive_utc().to_string(), + ":cidr": cidr_str, + ":task_group_id": task_group_id, + }).unwrap(); + } + + let two_hundred_millis = Duration::from_millis(500); + thread::sleep(two_hundred_millis); + } +});*/