Fix value validation

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"

snow-scanner/src/main.rs

@@ -2,51 +2,40 @@ 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};
+use diesel::deserialize::{self};
 use diesel::mysql::{Mysql, MysqlValue};
 use diesel::sql_types::Text;
 
 use diesel::r2d2::ConnectionManager;
 use diesel::r2d2::Pool;
+use worker::detection::{detect_scanner, get_dns_client, Scanners};
 
+use std::collections::HashMap;
 use std::io::Write;
 use std::path::PathBuf;
-use std::str::FromStr;
 use std::{env, fmt};
 use uuid::Uuid;
 
 use serde::{Deserialize, Deserializer, Serialize};
 
-use hickory_resolver::config::{NameServerConfigGroup, ResolverConfig, ResolverOpts};
-use hickory_resolver::{Name, Resolver};
-use std::net::IpAddr;
-use std::time::Duration;
-
 use diesel::serialize::IsNull;
 use diesel::{serialize, MysqlConnection};
 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<ConnectionManager<MysqlConnection>>;
 
-// Create alias for HMAC-SHA256
-// type HmacSha256 = Hmac<Sha256>;
-
-#[derive(Debug, Clone, Copy, FromSqlRow)]
-pub enum Scanners {
-    Stretchoid,
-    Binaryedge,
-    Censys,
-    InternetMeasurement,
-}
-
 trait IsStatic {
     fn is_static(self: &Self) -> bool;
 }
@@ -123,24 +112,6 @@ impl deserialize::FromSql<Text, Mysql> for Scanners {
     }
 }
 
-fn detect_scanner(ptr_result: &ResolvedResult) -> Result<Scanners, ()> {
-    match ptr_result.result {
-        Some(ref x)
-            if x.trim_to(2)
-                .eq_case(&Name::from_str("binaryedge.ninja.").expect("Should parse")) =>
-        {
-            Ok(Scanners::Binaryedge)
-        }
-        Some(ref x)
-            if x.trim_to(2)
-                .eq_case(&Name::from_str("stretchoid.com.").expect("Should parse")) =>
-        {
-            Ok(Scanners::Stretchoid)
-        }
-        _ => Err(()),
-    }
-}
-
 async fn handle_ip(pool: web::Data<DbPool>, ip: String) -> Result<Scanner, Option<ResolvedResult>> {
     let query_address = ip.parse().expect("To parse");
 
@@ -163,41 +134,12 @@ async fn handle_ip(pool: web::Data<DbPool>, ip: String) -> Result<Scanner, Optio
     let result = ptr_result.unwrap();
 
     match detect_scanner(&result) {
-        Ok(scanner_name) => {
+        Ok(Some(scanner_type)) => {
-            let ip_type = if ip.contains(':') { 6 } else { 4 };
-
             // use web::block to offload blocking Diesel queries without blocking server thread
             web::block(move || {
                 // note that obtaining a connection from the pool is also potentially blocking
                 let conn = &mut pool.get().unwrap();
-                let scanner_row_result = Scanner::find(ip.clone(), ip_type, conn);
+                match Scanner::find_or_new(query_address, scanner_type, result.result, conn) {
-                let scanner_row = match scanner_row_result {
-                    Ok(scanner_row) => scanner_row,
-                    Err(_) => return Err(None),
-                };
-
-                let scanner = if let Some(mut scanners) = scanner_row {
-                    scanners.last_seen_at = Some(Utc::now().naive_utc());
-                    scanners.last_checked_at = Some(Utc::now().naive_utc());
-                    scanners.updated_at = Some(Utc::now().naive_utc());
-                    scanners
-                } else {
-                    Scanner {
-                        ip: ip,
-                        ip_type: ip_type,
-                        scanner_name: scanner_name.clone(),
-                        ip_ptr: match result.result {
-                            Some(ptr) => Some(ptr.to_string()),
-                            None => None,
-                        },
-                        created_at: Utc::now().naive_utc(),
-                        updated_at: None,
-                        last_seen_at: None,
-                        last_checked_at: None,
-                    }
-                };
-
-                match scanner.save(conn) {
                     Ok(scanner) => Ok(scanner),
                     Err(_) => Err(None),
                 }
@@ -205,6 +147,7 @@ async fn handle_ip(pool: web::Data<DbPool>, ip: String) -> Result<Scanner, Optio
             .await
             .unwrap()
         }
+        Ok(None) => Err(None),
 
         Err(_) => Err(Some(result)),
     }
@@ -262,8 +205,8 @@ async fn handle_scan(pool: web::Data<DbPool>, params: web::Form<ScanParams>) ->
                 ended_at: None,
             };
             match scan_task.save(conn) {
-                Ok(_) => println!("Added {}", ip.to_string()),
+                Ok(_) => error!("Added {}", ip.to_string()),
-                Err(err) => eprintln!("Not added: {:?}", err),
+                Err(err) => error!("Not added: {:?}", err),
             }
         }
     })
@@ -329,18 +272,17 @@ impl<'de> Deserialize<'de> for SecurePath {
     where
         D: Deserializer<'de>,
     {
-        let s = <Vec<String>>::deserialize(deserializer)?;
+        let s = <String>::deserialize(deserializer)?;
-        let k: String = s[0].to_string();
         // A-Z a-z 0-9
         // . - _
-        if k.chars()
+        if s.chars()
             .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-' || c == '_')
         {
-            return Ok(SecurePath { data: k });
+            return Ok(SecurePath { data: s });
         }
         Err(serde::de::Error::custom(format!(
             "Invalid value: {}",
-            k.to_string()
+            s.to_string()
         )))
     }
 }
@@ -488,23 +430,6 @@ fn get_connection(database_url: &str) -> DbPool {
         .expect("Could not build connection pool")
 }
 
-fn get_dns_client() -> Resolver {
-    let server_ip = "1.1.1.1";
-
-    let server = NameServerConfigGroup::from_ips_clear(
-        &[IpAddr::from_str(server_ip).unwrap()],
-        53, // Port 53
-        true,
-    );
-
-    let config = ResolverConfig::from_parts(None, vec![], server);
-    let mut options = ResolverOpts::default();
-    options.timeout = Duration::from_secs(5);
-    options.attempts = 1; // One try
-
-    Resolver::new(config, options).unwrap()
-}
-
 fn plain_contents(data: String) -> HttpResponse {
     HttpResponse::Ok()
         .content_type(ContentType::plaintext())
@@ -533,16 +458,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 +491,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()),
+        Ok(names) => info!("Found {} Stretchoid scanners", names.len()),
-        Err(err) => eprintln!("Unable to get names: {}", err),
+        Err(err) => error!("Unable to get names: {}", err),
     };
 
     let server = HttpServer::new(move || {
@@ -582,108 +521,36 @@ 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 pool = get_connection(db_url.as_str());
+                        // note that obtaining a connection from the pool is also potentially blocking
+                        let conn = &mut pool.get().unwrap();
+                        let mut ws_server_handles = Server {
+                            clients: HashMap::new(),
+                            new_scanners: 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);
+                            ws_server_handles.commit(conn);
+                        }
+                    });
+                }
+                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! <b>{}</a>.", 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);
-        }
-    });*/

snow-scanner/src/models.rs

@@ -1,8 +1,11 @@
+use std::net::IpAddr;
+
 use crate::Scanners;
-use chrono::NaiveDateTime;
+use chrono::{NaiveDateTime, Utc};
 use diesel::dsl::insert_into;
 use diesel::prelude::*;
 use diesel::result::Error as DieselError;
+use hickory_resolver::Name;
 
 use crate::schema::scan_tasks::dsl::scan_tasks;
 use crate::schema::scanners::dsl::scanners;
@@ -22,6 +25,45 @@ pub struct Scanner {
 }
 
 impl Scanner {
+    pub fn find_or_new(
+        query_address: IpAddr,
+        scanner_name: Scanners,
+        ptr: Option<Name>,
+        conn: &mut MysqlConnection,
+    ) -> Result<Scanner, ()> {
+        let ip_type = if query_address.is_ipv6() { 6 } else { 4 };
+        let scanner_row_result = Scanner::find(query_address.to_string(), ip_type, conn);
+        let scanner_row = match scanner_row_result {
+            Ok(scanner_row) => scanner_row,
+            Err(_) => return Err(()),
+        };
+
+        let scanner = if let Some(mut scanner) = scanner_row {
+            scanner.last_seen_at = Some(Utc::now().naive_utc());
+            scanner.last_checked_at = Some(Utc::now().naive_utc());
+            scanner.updated_at = Some(Utc::now().naive_utc());
+            scanner
+        } else {
+            Scanner {
+                ip: query_address.to_string(),
+                ip_type: ip_type,
+                scanner_name: scanner_name.clone(),
+                ip_ptr: match ptr {
+                    Some(ptr) => Some(ptr.to_string()),
+                    None => None,
+                },
+                created_at: Utc::now().naive_utc(),
+                updated_at: None,
+                last_seen_at: None,
+                last_checked_at: None,
+            }
+        };
+        match scanner.save(conn) {
+            Ok(scanner) => Ok(scanner),
+            Err(_) => Err(()),
+        }
+    }
+
     pub fn find(
         ip_address: String,
         ip_type: u8,

snow-scanner/src/server.rs

@@ -0,0 +1,168 @@
+use cidr::IpCidr;
+use diesel::MysqlConnection;
+use hickory_resolver::Name;
+use log2::*;
+use std::{collections::HashMap, net::IpAddr, str::FromStr};
+use ws2::{Pod, WebSocket};
+
+use crate::{
+    worker::{
+        detection::detect_scanner_from_name,
+        modules::{Network, WorkerMessages},
+    },
+    DbPool, Scanner,
+};
+
+pub struct Server {
+    pub clients: HashMap<u32, Worker>,
+    pub new_scanners: HashMap<String, IpAddr>,
+}
+
+impl Server {
+    pub fn cleanup(&self, _: &ws2::Server) -> &Server {
+        // TODO: implement check not logged in
+        &self
+    }
+
+    pub fn commit(&mut self, conn: &mut MysqlConnection) -> &Server {
+        for (name, query_address) in self.new_scanners.clone() {
+            let scanner_name = Name::from_str(name.as_str()).unwrap();
+
+            match detect_scanner_from_name(&scanner_name) {
+                Ok(Some(scanner_type)) => {
+                    match Scanner::find_or_new(
+                        query_address,
+                        scanner_type,
+                        Some(scanner_name),
+                        conn,
+                    ) {
+                        Ok(scanner) => {
+                            // Got saved
+                            self.new_scanners.remove(&name);
+                            info!(
+                                "Saved {scanner_type}: {name} for {query_address}: {:?}",
+                                scanner.ip_ptr
+                            );
+                        }
+                        Err(err) => {
+                            error!("Unable to find or new {:?}", err);
+                        }
+                    };
+                }
+                Ok(None) => {}
+                Err(_) => {}
+            }
+        }
+        self
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Worker {
+    pub authenticated: bool,
+    pub login: Option<String>,
+}
+
+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 mut worker_reply: Option<WorkerMessages> = None;
+        let worker_request: WorkerMessages = msg.clone().into();
+
+        let result = match worker_request {
+            WorkerMessages::AuthenticateRequest { login } => {
+                if !worker.is_authenticated() {
+                    worker.authenticate(login);
+                    return Ok(());
+                } else {
+                    error!("Already authenticated: {ws}");
+                    return Ok(());
+                }
+            }
+            WorkerMessages::ScannerFoundResponse { name, address } => {
+                info!("Detected {name} for {address}");
+                self.new_scanners.insert(name, address);
+                Ok(())
+            }
+            WorkerMessages::GetWorkRequest {} => {
+                worker_reply = Some(WorkerMessages::DoWorkRequest {
+                    neworks: vec![Network(IpCidr::from_str("52.189.78.0/24")?)],
+                });
+                Ok(())
+            }
+            WorkerMessages::DoWorkRequest { .. } | WorkerMessages::Invalid { .. } => {
+                error!("Unable to understand: {msg}, {ws}");
+                // Unable to understand, close the connection
+                return ws.close();
+            } /*msg => {
+                  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
+    }
+}

snow-scanner/src/worker/detection.rs

@@ -0,0 +1,61 @@
+use std::net::IpAddr;
+use std::str::FromStr;
+use std::time::Duration;
+
+use diesel::deserialize::FromSqlRow;
+use dns_ptr_resolver::ResolvedResult;
+
+use hickory_resolver::config::{NameServerConfigGroup, ResolverConfig, ResolverOpts};
+use hickory_resolver::{Name, Resolver};
+
+#[derive(Debug, Clone, Copy, FromSqlRow)]
+pub enum Scanners {
+    Stretchoid,
+    Binaryedge,
+    Censys,
+    InternetMeasurement,
+}
+
+pub fn get_dns_client() -> Resolver {
+    let server_ip = "1.1.1.1";
+
+    let server = NameServerConfigGroup::from_ips_clear(
+        &[IpAddr::from_str(server_ip).unwrap()],
+        53, // Port 53
+        true,
+    );
+
+    let config = ResolverConfig::from_parts(None, vec![], server);
+    let mut options = ResolverOpts::default();
+    options.timeout = Duration::from_secs(5);
+    options.attempts = 1; // One try
+
+    Resolver::new(config, options).unwrap()
+}
+
+pub fn detect_scanner(ptr_result: &ResolvedResult) -> Result<Option<Scanners>, ()> {
+    match &ptr_result.result {
+        Some(name) => detect_scanner_from_name(&name),
+        None => Ok(None),
+    }
+}
+
+pub fn detect_scanner_from_name(name: &Name) -> Result<Option<Scanners>, ()> {
+    match name {
+        ref name
+            if name
+                .trim_to(2)
+                .eq_case(&Name::from_str("binaryedge.ninja.").expect("Should parse")) =>
+        {
+            Ok(Some(Scanners::Binaryedge))
+        }
+        ref name
+            if name
+                .trim_to(2)
+                .eq_case(&Name::from_str("stretchoid.com.").expect("Should parse")) =>
+        {
+            Ok(Some(Scanners::Stretchoid))
+        }
+        &_ => Ok(None),
+    }
+}

snow-scanner/src/worker/mod.rs

@@ -0,0 +1,2 @@
+pub mod detection;
+pub mod modules;

snow-scanner/src/worker/modules.rs

@@ -0,0 +1,99 @@
+use std::{net::IpAddr, str::FromStr};
+
+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")]
+    AuthenticateRequest { login: String },
+    #[serde(rename = "get_work")]
+    GetWorkRequest {},
+    #[serde(rename = "do_work")]
+    DoWorkRequest { neworks: Vec<Network> },
+    #[serde(rename = "scanner_found")]
+    ScannerFoundResponse { name: String, address: IpAddr },
+    #[serde(rename = "")]
+    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 {
+    fn to_string(&self) -> String {
+        serde_json::to_string(&self).expect("To serialize").into()
+    }
+}
+
+impl Into<WorkerMessages> for String {
+    fn into(self) -> WorkerMessages {
+        let req: Result<WorkerMessages, serde_json::Error> = serde_json::from_str(self.as_str());
+        match req {
+            Ok(d) => d,
+            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()
+            }
+        );
+    }
+}

snow-scanner/src/worker/worker.rs

@@ -0,0 +1,239 @@
+use std::{env, net::IpAddr};
+
+use chrono::{Duration, NaiveDateTime, Utc};
+use detection::detect_scanner;
+use dns_ptr_resolver::{get_ptr, ResolvedResult};
+use log2::*;
+use ws2::{Client, Pod, WebSocket};
+
+pub mod detection;
+pub mod modules;
+
+use crate::detection::get_dns_client;
+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 {
+    pub fn initial() -> Worker {
+        info!("New worker");
+        Worker {
+            authenticated: false,
+            tasks: vec![],
+            last_request_for_work: 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");
+            return self;
+        } else {
+            info!("Worker is now authenticated as {login}");
+        }
+        self.authenticated = true;
+        self
+    }
+
+    pub fn tick(&mut self, ws_client: &Client) -> &Worker {
+        let mut request: Option<WorkerMessages> = None;
+        if !self.is_authenticated() {
+            request = Some(WorkerMessages::AuthenticateRequest {
+                login: "williamdes".to_string(),
+            });
+        } else {
+            if self.last_request_for_work.is_none()
+                || (self.last_request_for_work.is_some()
+                    && Utc::now().naive_utc()
+                        > (self.last_request_for_work.unwrap() + Duration::minutes(10)))
+            {
+                request = Some(WorkerMessages::GetWorkRequest {});
+            }
+        }
+
+        // it has a request to send
+        if let Some(request) = request {
+            self.send_request(ws_client, request);
+        }
+        self
+    }
+
+    pub fn send_request(&mut self, ws_client: &Client, request: WorkerMessages) -> &Worker {
+        let msg_string: String = request.to_string();
+        match ws_client.send(msg_string) {
+            Ok(_) => {
+                match request {
+                    WorkerMessages::AuthenticateRequest { login } => {
+                        self.authenticated = true; // Anyway, it will kick us if this is not success
+                        info!("Logged in as: {login}")
+                    }
+                    WorkerMessages::GetWorkRequest {} => {
+                        self.last_request_for_work = Some(Utc::now().naive_utc());
+                        info!("Asked for work: {:?}", self.last_request_for_work)
+                    }
+                    msg => error!("No implemented: {:#?}", msg),
+                }
+            }
+            Err(err) => error!("Unable to send: {err}"),
+        }
+        self
+    }
+
+    pub fn receive_request(&mut self, ws: &WebSocket, server_request: WorkerMessages) -> &Worker {
+        match server_request {
+            WorkerMessages::DoWorkRequest { neworks } => {
+                info!("Should work on: {:?}", neworks);
+                for cidr in neworks {
+                    let cidr = cidr.0;
+                    info!("Picking up: {cidr}");
+                    info!("Range, from {} to {}", cidr.first(), cidr.last());
+                    let addresses = cidr.iter().addresses();
+                    let count = addresses.count();
+                    let mut current = 0;
+                    for addr in addresses {
+                        let client = get_dns_client();
+                        match get_ptr(addr, client) {
+                            Ok(result) => match detect_scanner(&result) {
+                                Ok(Some(scanner_name)) => {
+                                    info!("Detected {:?} for {addr}", scanner_name);
+                                    let request = WorkerMessages::ScannerFoundResponse {
+                                        name: result.result.unwrap().to_string(),
+                                        address: addr,
+                                    };
+                                    let msg_string: String = request.to_string();
+                                    match ws.send(msg_string) {
+                                        Ok(_) => {}
+                                        Err(err) => error!("Unable to send scanner result: {err}"),
+                                    }
+                                }
+                                Ok(None) => {}
+
+                                Err(err) => error!("Error detecting for {addr}: {:?}", err),
+                            },
+                            Err(err) => {
+                                //debug!("Error processing {addr}: {err}")
+                            }
+                        };
+
+                        current += 1;
+                    }
+                }
+            }
+            WorkerMessages::AuthenticateRequest { .. }
+            | WorkerMessages::ScannerFoundResponse { .. }
+            | WorkerMessages::GetWorkRequest {}
+            | WorkerMessages::Invalid { .. } => {
+                error!("Unable to understand message: {:?}", server_request);
+            }
+        }
+        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 {
+        let server_request: WorkerMessages = msg.clone().into();
+        self.receive_request(ws, server_request);
+        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(_) => {
+                        worker.tick(&ws_client);
+                    }
+                    Err(err) => error!("Processing error: {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 _ = 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();
+
+                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();
+        }
+    }
+});*/