Re-work crate management

snow-scanner/src/main.rs

@@ -27,29 +27,30 @@ use rocket_db_pools::{
     Connection, Pool,
 };
 
-use crate::worker::modules::{Network, WorkerMessages};
-
-use rocket_db_pools::diesel::mysql::{Mysql, MysqlValue};
-use rocket_db_pools::diesel::serialize::IsNull;
-use rocket_db_pools::diesel::sql_types::Text;
 use rocket_db_pools::diesel::MysqlPool;
-use rocket_db_pools::diesel::{deserialize, serialize};
 use rocket_db_pools::Database;
-
 use rocket_ws::WebSocket;
 use server::Server;
-use worker::detection::{detect_scanner, get_dns_client, validate_ip, Scanners};
+use weighted_rs::Weight;
 
+use snow_scanner_worker::detection::{
+    detect_scanner, get_dns_client, get_dns_server_config, validate_ip,
+};
+use snow_scanner_worker::modules::{Network, WorkerMessages};
+use snow_scanner_worker::scanners::IsStatic;
+use snow_scanner_worker::scanners::Scanners;
+use snow_scanner_worker::utils::get_dns_rr;
+
+use std::net::SocketAddr;
 use std::{
-    env, fmt,
+    env,
     net::IpAddr,
     ops::{Deref, DerefMut},
 };
-use std::{io::Write, net::SocketAddr};
 use std::{path::PathBuf, str::FromStr};
 use uuid::Uuid;
 
-use serde::{Deserialize, Deserializer, Serialize};
+use serde::{Deserialize, Serialize};
 
 use dns_ptr_resolver::{get_ptr, ResolvedResult};
 
@@ -57,7 +58,6 @@ pub mod event_bus;
 pub mod models;
 pub mod schema;
 pub mod server;
-pub mod worker;
 
 use crate::models::*;
 
@@ -99,20 +99,6 @@ impl<D: Database> DerefMut for DbConnection<D> {
     }
 }
 
-trait IsStatic {
-    fn is_static(self: &Self) -> bool;
-}
-
-impl IsStatic for Scanners {
-    fn is_static(self: &Self) -> bool {
-        match self {
-            Scanners::Censys => true,
-            Scanners::InternetMeasurement => true,
-            _ => false,
-        }
-    }
-}
-
 #[derive(serde::Deserialize, Clone)]
 struct SafeIpAddr {
     pub addr: IpAddr,
@@ -137,109 +123,12 @@ impl FromFormField<'_> for SafeIpAddr {
     }
 }
 
-impl FromParam<'_> for Scanners {
-    type Error = String;
-
-    fn from_param(param: &'_ str) -> Result<Self, Self::Error> {
-        match param {
-            "stretchoid" => Ok(Scanners::Stretchoid),
-            "binaryedge" => Ok(Scanners::Binaryedge),
-            "shadowserver" => Ok(Scanners::Shadowserver),
-            "stretchoid.txt" => Ok(Scanners::Stretchoid),
-            "binaryedge.txt" => Ok(Scanners::Binaryedge),
-            "shadowserver.txt" => Ok(Scanners::Shadowserver),
-            "censys.txt" => Ok(Scanners::Censys),
-            "internet-measurement.com.txt" => Ok(Scanners::InternetMeasurement),
-            v => Err(format!("Unknown value: {v}")),
-        }
-    }
-}
-
-impl<'de> Deserialize<'de> for Scanners {
-    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-    where
-        D: Deserializer<'de>,
-    {
-        let s = <Vec<String>>::deserialize(deserializer)?;
-        let k: &str = s[0].as_str();
-        match k {
-            "stretchoid" => Ok(Scanners::Stretchoid),
-            "binaryedge" => Ok(Scanners::Binaryedge),
-            "shadowserver" => Ok(Scanners::Shadowserver),
-            "stretchoid.txt" => Ok(Scanners::Stretchoid),
-            "binaryedge.txt" => Ok(Scanners::Binaryedge),
-            "shadowserver.txt" => Ok(Scanners::Shadowserver),
-            "censys.txt" => Ok(Scanners::Censys),
-            "internet-measurement.com.txt" => Ok(Scanners::InternetMeasurement),
-            v => Err(serde::de::Error::custom(format!(
-                "Unknown value: {}",
-                v.to_string()
-            ))),
-        }
-    }
-}
-
-impl fmt::Display for Scanners {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(
-            f,
-            "{}",
-            match self {
-                Self::Stretchoid => "stretchoid",
-                Self::Binaryedge => "binaryedge",
-                Self::Censys => "censys",
-                Self::InternetMeasurement => "internet-measurement.com",
-                Self::Shadowserver => "shadowserver",
-            }
-        )
-    }
-}
-
-impl serialize::ToSql<Text, Mysql> for Scanners {
-    fn to_sql(&self, out: &mut serialize::Output<Mysql>) -> serialize::Result {
-        match *self {
-            Self::Stretchoid => out.write_all(b"stretchoid")?,
-            Self::Binaryedge => out.write_all(b"binaryedge")?,
-            Self::Censys => out.write_all(b"censys")?,
-            Self::InternetMeasurement => out.write_all(b"internet-measurement.com")?,
-            Self::Shadowserver => out.write_all(b"shadowserver")?,
-        };
-
-        Ok(IsNull::No)
-    }
-}
-
-impl deserialize::FromSql<Text, Mysql> for Scanners {
-    fn from_sql(bytes: MysqlValue) -> deserialize::Result<Self> {
-        let value = <String as deserialize::FromSql<Text, Mysql>>::from_sql(bytes)?;
-        let value = &value as &str;
-        let value: Result<Scanners, String> = value.try_into();
-        match value {
-            Ok(d) => Ok(d),
-            Err(err) => Err(err.into()),
-        }
-    }
-}
-
-impl TryInto<Scanners> for &str {
-    type Error = String;
-
-    fn try_into(self) -> Result<Scanners, Self::Error> {
-        match self {
-            "stretchoid" => Ok(Scanners::Stretchoid),
-            "binaryedge" => Ok(Scanners::Binaryedge),
-            "internet-measurement.com" => Ok(Scanners::InternetMeasurement),
-            "shadowserver" => Ok(Scanners::Shadowserver),
-            value => Err(format!("Invalid value: {value}")),
-        }
-    }
-}
-
 async fn handle_ip(
     query_address: IpAddr,
 ) -> Result<(IpAddr, Option<Scanners>, ResolvedResult), ()> {
     let ptr_result: Result<ResolvedResult, ()> = std::thread::spawn(move || {
-        let client = get_dns_client();
+        let mut rr_dns_servers = get_dns_rr();
+        let client = get_dns_client(&get_dns_server_config(&rr_dns_servers.next().unwrap()));
         let ptr_result: ResolvedResult = if let Ok(res) = get_ptr(query_address, client) {
             res
         } else {