Re-work crate management

This commit is contained in:
2024-10-10 15:06:23 +02:00
parent 565e268d01
commit 1495950484
10 changed files with 348 additions and 198 deletions

View File

@ -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 {