From 7d900c5dd8fcd38bd09368d0e7e780e9984e2e15 Mon Sep 17 00:00:00 2001 From: ParrotXray Date: Sat, 23 May 2026 03:27:21 +0000 Subject: [PATCH] feat: adjust code with rustfmt --- common/src/model/event.rs | 1 - macros/src/log.rs | 4 +- mantis/src/core/ebpf/access_control.rs | 2 +- mantis/src/core/ebpf/mod.rs | 10 +- mantis/src/core/ebpf/service.rs | 2 +- mantis/src/core/ebpf/statistics.rs | 24 +++-- mantis/src/core/ebpf/xsk_manager.rs | 14 +-- mantis/src/core/infrastructure/app_config.rs | 2 +- mantis/src/core/infrastructure/app_db.rs | 16 +-- mantis/src/core/infrastructure/geoip.rs | 47 ++++----- mantis/src/core/infrastructure/health.rs | 58 ++++------- mantis/src/core/infrastructure/mod.rs | 26 +++-- mantis/src/core/system.rs | 12 +-- mantis/src/detection/fusion.rs | 24 ++++- mantis/src/detection/ml/aggregator.rs | 22 ++-- mantis/src/detection/ml/config_loader.rs | 20 ++-- mantis/src/detection/ml/engine.rs | 37 ++++--- mantis/src/detection/ml/feature_extractor.rs | 31 +++--- mantis/src/detection/ml/flow_tracker.rs | 49 +++++---- mantis/src/detection/ml/inference.rs | 30 +++--- mantis/src/detection/ml/mod.rs | 12 +-- mantis/src/detection/ml/model_loader.rs | 13 ++- mantis/src/detection/ml/traffic_logger.rs | 4 +- mantis/src/detection/mod.rs | 2 +- mantis/src/detection/suricata/engine.rs | 102 ++++++++++--------- mantis/src/detection/suricata/mod.rs | 2 +- mantis/src/detection/suricata/output.rs | 6 +- mantis/src/main.rs | 6 +- mantis/src/model/config.rs | 22 ++-- mantis/src/model/error/ebpf.rs | 1 - mantis/src/model/error/misc.rs | 2 +- mantis/src/model/error/ml.rs | 2 +- mantis/src/model/error/mod.rs | 4 +- mantis/src/model/geo_stats.rs | 2 +- mantis/src/model/ip_address.rs | 24 +---- mantis/src/model/log/ebpf.rs | 4 +- mantis/src/model/log/ml.rs | 2 +- mantis/src/model/log/mod.rs | 8 +- mantis/src/model/log/suricata.rs | 2 +- mantis/src/model/log/system.rs | 2 +- mantis/src/model/ml_detection.rs | 2 +- mantis/src/model/mod.rs | 4 +- mantis/src/utils/ip_address.rs | 9 +- mantis/src/utils/logging.rs | 6 +- mantis/src/utils/mod.rs | 8 +- mantis/src/utils/packet_parser.rs | 2 +- mantis/src/web/api/auth.rs | 26 +++-- mantis/src/web/api/control/access_control.rs | 14 ++- mantis/src/web/api/control/service.rs | 76 +++++++++++--- mantis/src/web/api/control/statistics.rs | 34 ++++--- mantis/src/web/api/default.rs | 15 +-- mantis/src/web/api/detection_alert.rs | 7 +- mantis/src/web/api/health.rs | 4 +- mantis/src/web/api/misc.rs | 4 +- mantis/src/web/middleware/auth.rs | 9 +- mantis/src/web/websocket/alert_websocket.rs | 2 +- mantis/src/web/websocket/flow_websocket.rs | 2 +- mantis/src/web/websocket/health_websocket.rs | 2 +- mantis/src/web/websocket/mod.rs | 2 +- 59 files changed, 488 insertions(+), 393 deletions(-) diff --git a/common/src/model/event.rs b/common/src/model/event.rs index e862aa4..5bd008d 100644 --- a/common/src/model/event.rs +++ b/common/src/model/event.rs @@ -109,7 +109,6 @@ impl Event { } } - #[repr(C, align(8))] #[derive(Debug, Clone)] pub struct IPv4Event { diff --git a/macros/src/log.rs b/macros/src/log.rs index ff56b3f..06a698f 100644 --- a/macros/src/log.rs +++ b/macros/src/log.rs @@ -1,7 +1,7 @@ use proc_macro::TokenStream; use quote::quote; use syn::parse::{Parse, ParseStream}; -use syn::{parse_macro_input, Expr, Token}; +use syn::{Expr, Token, parse_macro_input}; struct LogInput { error: Expr, @@ -62,5 +62,5 @@ pub fn log_impl(input: TokenStream) -> TokenStream { } } } - .into() + .into() } diff --git a/mantis/src/core/ebpf/access_control.rs b/mantis/src/core/ebpf/access_control.rs index 73df7d1..38b0206 100644 --- a/mantis/src/core/ebpf/access_control.rs +++ b/mantis/src/core/ebpf/access_control.rs @@ -8,8 +8,8 @@ use common::model::ip_address::{IPv4, IPv6, Port}; use tokio::sync::RwLock; use crate::model::direction::FlowDirection; -use crate::model::error::ebpf::EbpfError; use crate::model::error::Error; +use crate::model::error::ebpf::EbpfError; use crate::model::ip_address::NativeConvert; use crate::model::list_type::ListType; use crate::utils::ip_address::convert_ports_to_vec; diff --git a/mantis/src/core/ebpf/mod.rs b/mantis/src/core/ebpf/mod.rs index f3bccaa..b69e43e 100644 --- a/mantis/src/core/ebpf/mod.rs +++ b/mantis/src/core/ebpf/mod.rs @@ -17,8 +17,8 @@ use crate::core::ebpf::xsk_manager::XskManager; use crate::core::infrastructure::app_config::AppConfig; use crate::detection::ml::engine::Engine; use crate::detection::suricata::SuricataEngine; -use crate::model::error::system::SystemError; use crate::model::error::Error; +use crate::model::error::system::SystemError; pub struct EbpfServices { pub xsk_manager: Arc, @@ -29,11 +29,7 @@ pub struct EbpfServices { } impl EbpfServices { - pub fn new( - app_config: Arc, - ingress_ebpf: &mut Ebpf, - egress_ebpf: &mut Ebpf, - ) -> Result { + pub fn new(app_config: Arc, ingress_ebpf: &mut Ebpf, egress_ebpf: &mut Ebpf) -> Result { let xsk_manager = XskManager::new(app_config.clone(), ingress_ebpf, egress_ebpf)?; let access_control = AccessControl::new(ingress_ebpf)?; let service = Service::new(ingress_ebpf)?; @@ -70,4 +66,4 @@ impl EbpfServices { } } } -} \ No newline at end of file +} diff --git a/mantis/src/core/ebpf/service.rs b/mantis/src/core/ebpf/service.rs index 7afe782..0fe8568 100644 --- a/mantis/src/core/ebpf/service.rs +++ b/mantis/src/core/ebpf/service.rs @@ -8,8 +8,8 @@ use common::model::ip_address::{AddrPortV4, AddrPortV6, IPv4, IPv6}; use common::model::placeholder::PlaceHolder; use tokio::sync::RwLock; -use crate::model::error::ebpf::EbpfError; use crate::model::error::Error; +use crate::model::error::ebpf::EbpfError; use crate::model::ip_address::NativeConvert; pub struct Service { diff --git a/mantis/src/core/ebpf/statistics.rs b/mantis/src/core/ebpf/statistics.rs index 8d37afd..f783d55 100644 --- a/mantis/src/core/ebpf/statistics.rs +++ b/mantis/src/core/ebpf/statistics.rs @@ -160,11 +160,13 @@ impl Statistics { let map = self .ipv4_maps .get(&(direction, flow_direction, time_type)) - .ok_or_else(|| EbpfError::FlowMapKeyMissing( - format!("{:?}", direction), - format!("{:?}", flow_direction), - format!("{:?}", time_type), - ))?; + .ok_or_else(|| { + EbpfError::FlowMapKeyMissing( + format!("{:?}", direction), + format!("{:?}", flow_direction), + format!("{:?}", time_type), + ) + })?; let flow_data = map.write().await.get_map(); if let Some(ref geo_ip) = self.geo_ip { @@ -197,11 +199,13 @@ impl Statistics { let map = self .ipv6_maps .get(&(direction, flow_direction, time_type)) - .ok_or_else(|| EbpfError::FlowMapKeyMissing( - format!("{:?}", direction), - format!("{:?}", flow_direction), - format!("{:?}", time_type), - ))?; + .ok_or_else(|| { + EbpfError::FlowMapKeyMissing( + format!("{:?}", direction), + format!("{:?}", flow_direction), + format!("{:?}", time_type), + ) + })?; let flow_data = map.write().await.get_map(); if let Some(ref geo_ip) = self.geo_ip { diff --git a/mantis/src/core/ebpf/xsk_manager.rs b/mantis/src/core/ebpf/xsk_manager.rs index 3f67667..39bbb7f 100644 --- a/mantis/src/core/ebpf/xsk_manager.rs +++ b/mantis/src/core/ebpf/xsk_manager.rs @@ -6,9 +6,9 @@ use std::sync::Arc; use std::thread; use std::time::Duration; -use aya::maps::{MapData, XskMap}; use aya::Ebpf; -use crossbeam::channel::{bounded, Receiver, Sender}; +use aya::maps::{MapData, XskMap}; +use crossbeam::channel::{Receiver, Sender, bounded}; use crossbeam::queue::SegQueue; use macros::log; use parking_lot::Mutex; @@ -21,9 +21,9 @@ use crate::detection::ml::engine::Engine; use crate::detection::suricata::SuricataEngine; use crate::model::config::Config; use crate::model::direction::Direction; +use crate::model::error::Error; use crate::model::error::ebpf::EbpfError; use crate::model::error::system::SystemError; -use crate::model::error::Error; use crate::model::log::ebpf::EbpfLog; use crate::utils::cpu_affinity::set_cpu_affinity; @@ -210,7 +210,6 @@ impl XskPair { set_cpu_affinity((start + self.queue_id % num_cores) as usize); } - let mut shutdown_rx = Some(shutdown_rx); let mut idle_count: u32 = 0; let mut last_cleanup = std::time::Instant::now(); @@ -286,10 +285,7 @@ impl XskPair { Ok(nb_completed) } - fn process_rx_queue( - &mut self, - forward_tx: &Sender>, - ) -> Result { + fn process_rx_queue(&mut self, forward_tx: &Sender>) -> Result { let mut rx_descs = vec![FrameDesc::default(); 64]; let rx_count = unsafe { self.rx.consume(&mut rx_descs) }; @@ -394,4 +390,4 @@ impl XskPair { Ok(nb_submitted) } -} \ No newline at end of file +} diff --git a/mantis/src/core/infrastructure/app_config.rs b/mantis/src/core/infrastructure/app_config.rs index 7f67ee0..a01e003 100644 --- a/mantis/src/core/infrastructure/app_config.rs +++ b/mantis/src/core/infrastructure/app_config.rs @@ -2,8 +2,8 @@ use std::fs; use std::ops::Deref; use crate::model::config::{Config, ConfigTable}; -use crate::model::error::system::SystemError; use crate::model::error::Error; +use crate::model::error::system::SystemError; pub struct AppConfig { pub config: Config, diff --git a/mantis/src/core/infrastructure/app_db.rs b/mantis/src/core/infrastructure/app_db.rs index 82cca75..41dedb1 100644 --- a/mantis/src/core/infrastructure/app_db.rs +++ b/mantis/src/core/infrastructure/app_db.rs @@ -1,13 +1,13 @@ use std::path::Path; use std::sync::Mutex; -use argon2::password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString}; use argon2::Argon2; +use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng}; use macros::log; -use rusqlite::{params, Connection}; +use rusqlite::{Connection, params}; -use crate::model::error::auth::AuthError; use crate::model::error::Error; +use crate::model::error::auth::AuthError; use crate::model::log::auth::AuthLog; pub struct Account { @@ -26,12 +26,10 @@ impl AppDB { let path = path.as_ref(); if let Some(parent) = Path::new(path).parent() { - std::fs::create_dir_all(parent) - .map_err(|e| AuthError::DBError { msg: e.to_string() })?; + std::fs::create_dir_all(parent).map_err(|e| AuthError::DBError { msg: e.to_string() })?; } - let conn = Connection::open(path) - .map_err(|e| AuthError::DBError { msg: e.to_string() })?; + let conn = Connection::open(path).map_err(|e| AuthError::DBError { msg: e.to_string() })?; // Must be the first statement on the connection to unlock the encrypted DB. conn.execute_batch(&format!("PRAGMA key = '{}';", key.replace('\'', "''"))) @@ -55,7 +53,9 @@ impl AppDB { ) .map_err(|e| AuthError::DBError { msg: e.to_string() })?; - log!(AuthLog::DbInitialized { path: path.to_string_lossy().to_string() }); + log!(AuthLog::DbInitialized { + path: path.to_string_lossy().to_string() + }); Ok(Self { conn: Mutex::new(conn) }) } diff --git a/mantis/src/core/infrastructure/geoip.rs b/mantis/src/core/infrastructure/geoip.rs index 1a59d41..8ce73d4 100644 --- a/mantis/src/core/infrastructure/geoip.rs +++ b/mantis/src/core/infrastructure/geoip.rs @@ -1,11 +1,11 @@ use std::net::IpAddr; +use std::num::NonZeroUsize; use std::path::{Path, PathBuf}; use std::sync::Arc; -use maxminddb::{geoip2, MaxMindDbError, Reader}; -use tokio::sync::{RwLock, Semaphore}; use lru::LruCache; -use std::num::NonZeroUsize; +use maxminddb::{MaxMindDbError, Reader, geoip2}; +use tokio::sync::{RwLock, Semaphore}; use tokio::task; use crate::model::geo_stats::GeoLocation; @@ -25,13 +25,10 @@ impl GeoIpService { Self::with_cache_size(db_path, 10000) } - pub fn with_cache_size>( - db_path: P, - cache_size: usize, - ) -> Result { + pub fn with_cache_size>(db_path: P, cache_size: usize) -> Result { let reader = Reader::open_readfile(db_path)?; - let cache_capacity = NonZeroUsize::new(cache_size) - .unwrap_or_else(|| NonZeroUsize::new(10000).expect("10000 is non-zero")); + let cache_capacity = + NonZeroUsize::new(cache_size).unwrap_or_else(|| NonZeroUsize::new(10000).expect("10000 is non-zero")); Ok(Self { reader: Arc::new(reader), @@ -59,7 +56,11 @@ impl GeoIpService { } } - let permit = self.lookup_sem.clone().acquire_owned().await + let permit = self + .lookup_sem + .clone() + .acquire_owned() + .await .map_err(|_| MaxMindDbError::InvalidDatabase { message: "GeoIP semaphore closed".to_string(), offset: None, @@ -70,11 +71,11 @@ impl GeoIpService { let _permit = permit; Self::lookup_from_db_blocking(&reader, ip) }) - .await - .map_err(|e| MaxMindDbError::InvalidDatabase { - message: format!("Task join error: {}", e), - offset: None, - })??; + .await + .map_err(|e| MaxMindDbError::InvalidDatabase { + message: format!("Task join error: {}", e), + offset: None, + })??; { let mut cache = self.cache.write().await; @@ -84,22 +85,16 @@ impl GeoIpService { Ok(result) } - fn lookup_from_db_blocking( - reader: &Reader>, - ip: IpAddr, - ) -> Result, MaxMindDbError> { + fn lookup_from_db_blocking(reader: &Reader>, ip: IpAddr) -> Result, MaxMindDbError> { let lookup_result = reader.lookup(ip)?; let city_option: Option = lookup_result.decode()?; Ok(city_option.map(|city| { - let country_name = city.country.names.english - .map(|s| s.to_string()); + let country_name = city.country.names.english.map(|s| s.to_string()); - let country_code = city.country.iso_code - .map(|s| s.to_string()); + let country_code = city.country.iso_code.map(|s| s.to_string()); - let city_name = city.city.names.english - .map(|s| s.to_string()); + let city_name = city.city.names.english.map(|s| s.to_string()); let latitude = city.location.latitude.or(Some(0.0)); let longitude = city.location.longitude.or(Some(0.0)); @@ -120,4 +115,4 @@ impl GeoIpService { let cache = self.cache.read().await; (cache.len(), cache.cap().get()) } -} \ No newline at end of file +} diff --git a/mantis/src/core/infrastructure/health.rs b/mantis/src/core/infrastructure/health.rs index ce4b33d..3b31373 100644 --- a/mantis/src/core/infrastructure/health.rs +++ b/mantis/src/core/infrastructure/health.rs @@ -1,25 +1,18 @@ use std::sync::Arc; use std::time::Duration; -use sysinfo::{Components, Networks, System}; -use tokio::sync::{broadcast, oneshot, RwLock}; -use tokio::time::interval; use macros::log; +use sysinfo::{Components, Networks, System}; +use tokio::sync::{RwLock, broadcast, oneshot}; +use tokio::time::interval; use crate::core::infrastructure::app_config::AppConfig; -use crate::model::log::health::Health; use crate::model::error::Error; use crate::model::health::{ - ConfiguredNetworkStats, - CpuCoreInfo, - CpuDetails, - LoadAverage, - MemoryUsage, - NetworkStats, - SystemHealthMetrics, - SystemHealthStatus, - SystemInfo + ConfiguredNetworkStats, CpuCoreInfo, CpuDetails, LoadAverage, MemoryUsage, NetworkStats, SystemHealthMetrics, + SystemHealthStatus, SystemInfo, }; +use crate::model::log::health::Health; pub struct SystemHealth { system: RwLock, @@ -30,7 +23,6 @@ pub struct SystemHealth { egress_interface: String, } - impl SystemHealth { pub fn new(config: Arc) -> Result { let (broadcast_tx, _) = broadcast::channel(100); @@ -125,11 +117,7 @@ impl SystemHealth { swap_used: system.used_swap(), }; - let network_stats = Self::collect_configured_network_stats( - networks, - ingress_interface, - egress_interface, - ); + let network_stats = Self::collect_configured_network_stats(networks, ingress_interface, egress_interface); let load_average = System::load_average(); let load_average = if load_average.one != 0.0 || load_average.five != 0.0 || load_average.fifteen != 0.0 { @@ -229,16 +217,19 @@ impl SystemHealth { let egress = create_network_stats(egress_interface); if ingress.is_none() { - log!(Health::InterfaceNotFound("Ingress".to_string(), ingress_interface.to_string())); + log!(Health::InterfaceNotFound( + "Ingress".to_string(), + ingress_interface.to_string() + )); } if egress.is_none() { - log!(Health::InterfaceNotFound("Egress".to_string(), egress_interface.to_string())); + log!(Health::InterfaceNotFound( + "Egress".to_string(), + egress_interface.to_string() + )); } - ConfiguredNetworkStats { - ingress, - egress, - } + ConfiguredNetworkStats { ingress, egress } } pub async fn get_current_metrics(&self) -> SystemHealthMetrics { @@ -290,22 +281,17 @@ impl SystemHealth { metrics.memory_usage.usage_percent )); } else if metrics.memory_usage.usage_percent > 80.0 { - status.warnings.push(format!( - "High memory usage: {:.1}%", - metrics.memory_usage.usage_percent - )); + status + .warnings + .push(format!("High memory usage: {:.1}%", metrics.memory_usage.usage_percent)); } if let Some(temp) = metrics.temperature { if temp > 80.0 { status.overall_healthy = false; - status - .issues - .push(format!("High CPU temperature: {:.1}°C", temp)); + status.issues.push(format!("High CPU temperature: {:.1}°C", temp)); } else if temp > 70.0 { - status - .warnings - .push(format!("Elevated CPU temperature: {:.1}°C", temp)); + status.warnings.push(format!("Elevated CPU temperature: {:.1}°C", temp)); } } @@ -320,4 +306,4 @@ impl SystemHealth { status } -} \ No newline at end of file +} diff --git a/mantis/src/core/infrastructure/mod.rs b/mantis/src/core/infrastructure/mod.rs index ed38c55..4227c78 100644 --- a/mantis/src/core/infrastructure/mod.rs +++ b/mantis/src/core/infrastructure/mod.rs @@ -1,8 +1,8 @@ pub mod app_config; pub mod app_db; pub mod detection_alert; -pub mod health; pub mod geoip; +pub mod health; use std::path::PathBuf; use std::sync::Arc; @@ -24,8 +24,8 @@ use crate::detection::ml::feature_extractor::FlowFeatures; use crate::detection::ml::model_loader::MLModels; use crate::detection::ml::traffic_logger::TrafficLogger; use crate::detection::suricata::SuricataEngine; -use crate::model::error::misc::MiscError; use crate::model::error::Error; +use crate::model::error::misc::MiscError; use crate::model::error::system::SystemError; use crate::model::log::system::SystemLog; @@ -56,12 +56,21 @@ impl AppServices { let traffic_logger = if app_config.traffic_logging_mode { let dir = PathBuf::from(env!("CSV_RECORD_PATH")); let basename = app_config.traffic_log_csv_path.trim_end_matches(".csv").to_string(); - let mut header = vec!["Source IP".to_string(), "Destination IP".to_string(), "Timestamp".to_string()]; + let mut header = vec![ + "Source IP".to_string(), + "Destination IP".to_string(), + "Timestamp".to_string(), + ]; header.extend(FlowFeatures::all_feature_names_owned()); header.push("Label".to_string()); let logger = TrafficLogger::new(&dir, &basename, header) .map_err(|e| MiscError::TrafficLogCreateError(dir.display().to_string(), e.to_string()))?; - log!(SystemLog::TrafficLoggingEnabled(format!("{}/{}-{}.csv", dir.display(), basename, Local::now().format("%Y-%m-%d")))); + log!(SystemLog::TrafficLoggingEnabled(format!( + "{}/{}-{}.csv", + dir.display(), + basename, + Local::now().format("%Y-%m-%d") + ))); Some(Arc::new(logger)) } else { None @@ -83,7 +92,12 @@ impl AppServices { let suricata_engine = if let Some(ref sc) = app_config.suricata { let rule_path = PathBuf::from(env!("RULE_PATH")); let eve_socket = PathBuf::from(env!("RULE_EVE_PATH")); - Some(SuricataEngine::start(sc, &rule_path, &eve_socket, fusion_engine.clone())?) + Some(SuricataEngine::start( + sc, + &rule_path, + &eve_socket, + fusion_engine.clone(), + )?) } else { None }; @@ -129,4 +143,4 @@ impl AppServices { } } } -} \ No newline at end of file +} diff --git a/mantis/src/core/system.rs b/mantis/src/core/system.rs index 1e59c7c..b72a26c 100644 --- a/mantis/src/core/system.rs +++ b/mantis/src/core/system.rs @@ -1,10 +1,10 @@ use std::sync::Arc; +use axum::Router; +use aya::Ebpf; use aya::maps::{MapData, ProgramArray}; use aya::programs::{Xdp, XdpFlags}; -use aya::Ebpf; use aya_log::EbpfLogger; -use axum::Router; use common::define::program_array::*; use macros::log; use tokio::net::TcpListener; @@ -12,18 +12,18 @@ use tower_http::cors::CorsLayer; use crate::core::app_state::AppState; use crate::core::ebpf::EbpfServices; -use crate::core::infrastructure::app_config::AppConfig; use crate::core::infrastructure::AppServices; +use crate::core::infrastructure::app_config::AppConfig; use crate::detection::ml::config_loader::InferenceConfig; +use crate::model::error::Error; use crate::model::error::ebpf::EbpfError; use crate::model::error::http::HttpError; use crate::model::error::misc::MiscError; -use crate::model::error::Error; use crate::model::log::ml::MLLog; use crate::model::log::system::SystemLog; use crate::utils::logging::Logging; -use crate::web::api::{auth, control, detection_alert, health, misc}; use crate::web::api::default::default_route; +use crate::web::api::{auth, control, detection_alert, health, misc}; pub struct System { pub app_config: Arc, @@ -243,4 +243,4 @@ impl System { } Ok(()) } -} \ No newline at end of file +} diff --git a/mantis/src/detection/fusion.rs b/mantis/src/detection/fusion.rs index b88b988..13e84a1 100644 --- a/mantis/src/detection/fusion.rs +++ b/mantis/src/detection/fusion.rs @@ -14,7 +14,11 @@ pub enum FusionMode { impl FusionMode { pub fn from_str(s: &str) -> Self { - if s.eq_ignore_ascii_case("and") { FusionMode::And } else { FusionMode::Or } + if s.eq_ignore_ascii_case("and") { + FusionMode::And + } else { + FusionMode::Or + } } } @@ -60,7 +64,14 @@ impl FusionEngine { return; } self.alert.broadcast(UnifiedAlert::from_ml(result)); - map.insert(key, FusionState { ml: Some(result.clone()), rule: None, created_at: now }); + map.insert( + key, + FusionState { + ml: Some(result.clone()), + rule: None, + created_at: now, + }, + ); } FusionMode::And => { let maybe_rule = { @@ -100,7 +111,14 @@ impl FusionEngine { return; } self.alert.broadcast(UnifiedAlert::from_rule(m)); - map.insert(key, FusionState { ml: None, rule: Some(m.clone()), created_at: now }); + map.insert( + key, + FusionState { + ml: None, + rule: Some(m.clone()), + created_at: now, + }, + ); } FusionMode::And => { let maybe_ml = { diff --git a/mantis/src/detection/ml/aggregator.rs b/mantis/src/detection/ml/aggregator.rs index d1ee27b..cadccbe 100644 --- a/mantis/src/detection/ml/aggregator.rs +++ b/mantis/src/detection/ml/aggregator.rs @@ -5,7 +5,7 @@ use crate::model::ml_detection::FlowKey; // L2 thresholds: anomalous flows per src_ip within the aggregation window const FLOOD_THRESHOLD: usize = 10; // total anomalous flows -> FLOOD -const SCAN_THRESHOLD: usize = 5; // distinct dst_ports -> SCAN +const SCAN_THRESHOLD: usize = 5; // distinct dst_ports -> SCAN struct SrcIpState { events: Vec<(Instant, u16)>, // (time, dst_port) @@ -39,8 +39,7 @@ impl AttackAggregator { detections.push((now, score)); if detections.len() >= self.min_detections { - let avg_score: f32 = - detections.iter().map(|(_, s)| s).sum::() / detections.len() as f32; + let avg_score: f32 = detections.iter().map(|(_, s)| s).sum::() / detections.len() as f32; return avg_score > threshold * self.alert_threshold_multiplier; } @@ -52,10 +51,13 @@ impl AttackAggregator { let now = Instant::now(); let window = self.window_duration; - let state = self.src_ip_states.entry(src_ip.to_string()).or_insert_with(|| SrcIpState { - events: Vec::new(), - last_alert: None, - }); + let state = self + .src_ip_states + .entry(src_ip.to_string()) + .or_insert_with(|| SrcIpState { + events: Vec::new(), + last_alert: None, + }); state.events.retain(|(t, _)| now.duration_since(*t) < window); state.events.push((now, dst_port)); @@ -89,7 +91,9 @@ impl AttackAggregator { !detections.is_empty() }); self.src_ip_states.retain(|_, state| { - state.events.retain(|(t, _)| now.duration_since(*t) < self.window_duration); + state + .events + .retain(|(t, _)| now.duration_since(*t) < self.window_duration); !state.events.is_empty() }); } @@ -97,4 +101,4 @@ impl AttackAggregator { pub fn tracked_flows(&self) -> usize { self.detections.len() } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/config_loader.rs b/mantis/src/detection/ml/config_loader.rs index b2e6646..57bf6a3 100644 --- a/mantis/src/detection/ml/config_loader.rs +++ b/mantis/src/detection/ml/config_loader.rs @@ -1,9 +1,10 @@ use std::collections::HashMap; use std::fs; use std::path::PathBuf; -use serde::{Deserialize, Serialize}; use macros::log; +use serde::{Deserialize, Serialize}; + use crate::model::error::ml::MLError; use crate::model::ml_detection::ClipParams; @@ -25,15 +26,18 @@ pub struct InferenceConfig { impl InferenceConfig { pub fn load_file(file: &str, method: &str) -> Result { let path = PathBuf::from(env!("ARTIFACTCS_PATH")).join(file); - let content = fs::read_to_string(&path) - .map_err(|_| MLError::ConfigLoadFailed { path: path.clone() })?; - let mut config: InferenceConfig = serde_json::from_str(&content) - .map_err(|e| MLError::ConfigParseFailed { reason: e.to_string() })?; + let content = fs::read_to_string(&path).map_err(|_| MLError::ConfigLoadFailed { path: path.clone() })?; + let mut config: InferenceConfig = + serde_json::from_str(&content).map_err(|e| MLError::ConfigParseFailed { reason: e.to_string() })?; // Validate the method exists at startup — fail fast rather than at inference time. if !config.ae_thresholds.contains_key(method) { - log!(MLError::ThresholdMethodNotFound { method: method.to_string() }); - return Err(MLError::ThresholdMethodNotFound { method: method.to_string() }); + log!(MLError::ThresholdMethodNotFound { + method: method.to_string() + }); + return Err(MLError::ThresholdMethodNotFound { + method: method.to_string(), + }); } config.ae_threshold_method = method.to_string(); @@ -43,4 +47,4 @@ impl InferenceConfig { pub fn num_ae_features(&self) -> usize { self.ae_feature_names.len() } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/engine.rs b/mantis/src/detection/ml/engine.rs index 49d9c00..1cad052 100644 --- a/mantis/src/detection/ml/engine.rs +++ b/mantis/src/detection/ml/engine.rs @@ -12,7 +12,6 @@ use super::flow_tracker::FlowTracker; use super::inference::Inference; use super::model_loader::MLModels; use super::traffic_logger::TrafficLogger; - use crate::detection::fusion::FusionEngine; use crate::model::error::ml::MLError; use crate::model::log::ml::MLLog; @@ -81,9 +80,9 @@ impl Engine { loop { tokio::select! { - _ = &mut shutdown_rx => break, - _ = ticker.tick() => {} - } + _ = &mut shutdown_rx => break, + _ = ticker.tick() => {} + } let (total_flows, flows, active_ips) = { let Ok(mut t) = self.tracker.lock() else { @@ -97,10 +96,8 @@ impl Engine { // (which is always empty). This preserves per-src_ip inference buffers // across consecutive cycles so the LSTM window can fill up over time. // A src_ip absent from this cycle loses its buffer on the next cleanup. - let active_ips: std::collections::HashSet = flows - .iter() - .map(|f| f.flow_key.src_ip.clone()) - .collect(); + let active_ips: std::collections::HashSet = + flows.iter().map(|f| f.flow_key.src_ip.clone()).collect(); (total_flows, flows, active_ips) }; @@ -122,7 +119,9 @@ impl Engine { let mut batch = flows[..flows.len().min(self.batch_size)].to_vec(); batch.sort_by(|a, b| { - a.flow_key.src_ip.cmp(&b.flow_key.src_ip) + a.flow_key + .src_ip + .cmp(&b.flow_key.src_ip) .then_with(|| a.start_time_us.cmp(&b.start_time_us)) }); @@ -131,7 +130,9 @@ impl Engine { let pipeline = Arc::clone(&self.inference_pipeline); let ml_cpu = self.ml_cpu; let mut handle = tokio::task::spawn_blocking(move || { - let cpu = ml_cpu.map(|c| c as usize).unwrap_or_else(|| num_cpus().saturating_sub(1)); + let cpu = ml_cpu + .map(|c| c as usize) + .unwrap_or_else(|| num_cpus().saturating_sub(1)); set_cpu_affinity(cpu); pipeline.infer_batch(&batch) }); @@ -166,11 +167,8 @@ impl Engine { for result in &results { if result.is_attack { // L1: full 5-tuple aggregation for persistent same-port attacks - let should_alert = aggregator.should_alert( - &result.flow_key_raw, - result.ae_score, - result.threshold, - ); + let should_alert = + aggregator.should_alert(&result.flow_key_raw, result.ae_score, result.threshold); if should_alert { log!(MLLog::ThreatDetected( format!("{:?}", result.direction), @@ -185,10 +183,9 @@ impl Engine { // L2: src_ip-level scan/flood detection (skipped if L1 already fired) if !alerted_src_ips.contains(&result.flow_key_raw.src_ip) { - if let Some(attack_type) = aggregator.should_alert_src_ip( - &result.flow_key_raw.src_ip, - result.flow_key_raw.dst_port, - ) { + if let Some(attack_type) = aggregator + .should_alert_src_ip(&result.flow_key_raw.src_ip, result.flow_key_raw.dst_port) + { let mut l2_result = result.clone(); l2_result.attack_type = Some(attack_type.to_string()); log!(MLLog::ThreatDetected( @@ -235,4 +232,4 @@ impl Engine { }; EngineStats { active_flows } } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/feature_extractor.rs b/mantis/src/detection/ml/feature_extractor.rs index 3b2e236..2843c74 100644 --- a/mantis/src/detection/ml/feature_extractor.rs +++ b/mantis/src/detection/ml/feature_extractor.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; -use chrono::{Utc, TimeZone}; + +use chrono::{TimeZone, Utc}; use super::flow_tracker::FlowData; use crate::model::ml_detection::{ClipParams, PacketData}; @@ -67,7 +68,9 @@ impl FlowFeatures { let bwd_urg = flow.bwd_packets.iter().filter(|p| p.flags.urg).count() as f64; let all_lengths: Vec = flow - .fwd_packets.iter().chain(flow.bwd_packets.iter()) + .fwd_packets + .iter() + .chain(flow.bwd_packets.iter()) .map(|p| p.payload_length as f64) .collect(); let (max_len, min_len, mean_len, std_len) = compute_stats(&all_lengths); @@ -75,11 +78,7 @@ impl FlowFeatures { let fwd_bulk = &flow.fwd_bulk_state; let bwd_bulk = &flow.bwd_bulk_state; - let fwd_header_sizes: Vec = flow - .fwd_packets - .iter() - .map(|p| p.header_length as f64) - .collect(); + let fwd_header_sizes: Vec = flow.fwd_packets.iter().map(|p| p.header_length as f64).collect(); let (active_max, active_min, active_mean, active_std) = compute_stats(&flow.active_periods.iter().map(|&x| x as f64).collect::>()); @@ -145,14 +144,22 @@ impl FlowFeatures { "Avg Fwd Segment Size" | "avg_fwd_seg_size" => safe_div(flow.fwd_total_bytes as f64, fwd_count), "Avg Bwd Segment Size" | "avg_bwd_seg_size" => safe_div(flow.bwd_total_bytes as f64, bwd_count), "Fwd Header Length.1" | "fwd_header_length_1" => flow.fwd_header_bytes as f64, - "Fwd Avg Bytes/Bulk" | "fwd_avg_bytes_bulk" => safe_div(fwd_bulk.total_bytes as f64, fwd_bulk.bulk_count as f64), - "Fwd Avg Packets/Bulk" | "fwd_avg_pkts_bulk" => safe_div(fwd_bulk.total_packets as f64, fwd_bulk.bulk_count as f64), + "Fwd Avg Bytes/Bulk" | "fwd_avg_bytes_bulk" => { + safe_div(fwd_bulk.total_bytes as f64, fwd_bulk.bulk_count as f64) + } + "Fwd Avg Packets/Bulk" | "fwd_avg_pkts_bulk" => { + safe_div(fwd_bulk.total_packets as f64, fwd_bulk.bulk_count as f64) + } "Fwd Avg Bulk Rate" | "fwd_avg_bulk_rate" => safe_div( fwd_bulk.total_bytes as f64, fwd_bulk.total_duration_us as f64 / 1_000_000.0, ), - "Bwd Avg Bytes/Bulk" | "bwd_avg_bytes_bulk" => safe_div(bwd_bulk.total_bytes as f64, bwd_bulk.bulk_count as f64), - "Bwd Avg Packets/Bulk" | "bwd_avg_pkts_bulk" => safe_div(bwd_bulk.total_packets as f64, bwd_bulk.bulk_count as f64), + "Bwd Avg Bytes/Bulk" | "bwd_avg_bytes_bulk" => { + safe_div(bwd_bulk.total_bytes as f64, bwd_bulk.bulk_count as f64) + } + "Bwd Avg Packets/Bulk" | "bwd_avg_pkts_bulk" => { + safe_div(bwd_bulk.total_packets as f64, bwd_bulk.bulk_count as f64) + } "Bwd Avg Bulk Rate" | "bwd_avg_bulk_rate" => safe_div( bwd_bulk.total_bytes as f64, bwd_bulk.total_duration_us as f64 / 1_000_000.0, @@ -358,4 +365,4 @@ fn compute_flow_iats(fwd_packets: &[PacketData], bwd_packets: &[PacketData]) -> .windows(2) .map(|w| (w[1].timestamp_us - w[0].timestamp_us) as f64) .collect() -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/flow_tracker.rs b/mantis/src/detection/ml/flow_tracker.rs index 543c61d..bf76fad 100644 --- a/mantis/src/detection/ml/flow_tracker.rs +++ b/mantis/src/detection/ml/flow_tracker.rs @@ -103,13 +103,27 @@ impl FlowData { self.bwd_fin_seen = true; } } - if packet.tcp_flags().syn { self.syn_count += 1; } - if packet.tcp_flags().rst { self.rst_count += 1; } - if packet.tcp_flags().psh { self.psh_count += 1; } - if packet.tcp_flags().ack { self.ack_count += 1; } - if packet.tcp_flags().urg { self.urg_count += 1; } - if packet.tcp_flags().cwr { self.cwe_count += 1; } - if packet.tcp_flags().ece { self.ece_count += 1; } + if packet.tcp_flags().syn { + self.syn_count += 1; + } + if packet.tcp_flags().rst { + self.rst_count += 1; + } + if packet.tcp_flags().psh { + self.psh_count += 1; + } + if packet.tcp_flags().ack { + self.ack_count += 1; + } + if packet.tcp_flags().urg { + self.urg_count += 1; + } + if packet.tcp_flags().cwr { + self.cwe_count += 1; + } + if packet.tcp_flags().ece { + self.ece_count += 1; + } let iat = packet.timestamp_us().saturating_sub(self.last_packet_time); const IDLE_THRESHOLD_US: u64 = 1_000_000; @@ -166,9 +180,7 @@ impl FlowData { if packet.payload_length > 0 { // Idle break: discard helper, start fresh - if bulk_state.in_bulk - && packet.timestamp_us.saturating_sub(bulk_state.last_bulk_packet_us) > BULK_IDLE_US - { + if bulk_state.in_bulk && packet.timestamp_us.saturating_sub(bulk_state.last_bulk_packet_us) > BULK_IDLE_US { bulk_state.in_bulk = false; bulk_state.last_bulk_bytes = 0; bulk_state.last_bulk_packets = 0; @@ -193,14 +205,12 @@ impl FlowData { bulk_state.bulk_count += 1; bulk_state.total_packets += 4; bulk_state.total_bytes += bulk_state.last_bulk_bytes; - bulk_state.total_duration_us += packet.timestamp_us - .saturating_sub(bulk_state.last_bulk_start_us); + bulk_state.total_duration_us += packet.timestamp_us.saturating_sub(bulk_state.last_bulk_start_us); } else if bulk_state.last_bulk_packets > 4 { // Each subsequent packet adds incrementally bulk_state.total_packets += 1; bulk_state.total_bytes += packet.length as u64; - bulk_state.total_duration_us += packet.timestamp_us - .saturating_sub(prev_us); + bulk_state.total_duration_us += packet.timestamp_us.saturating_sub(prev_us); } } } else { @@ -279,7 +289,8 @@ impl FlowTracker { let initiator_direction = if is_forward { direction } else { direction.flip() }; - let flow = self.flows + let flow = self + .flows .entry(actual_key.clone()) .or_insert_with(|| FlowData::new(actual_key, &packet, initiator_direction)); @@ -305,10 +316,8 @@ impl FlowTracker { .duration_since(time::UNIX_EPOCH) .map(|d| d.as_micros() as u64) .unwrap_or(0); - self.flows.retain(|_, flow| { - !flow.is_finished() - && now.saturating_sub(flow.last_time_us) < max_age_us - }); + self.flows + .retain(|_, flow| !flow.is_finished() && now.saturating_sub(flow.last_time_us) < max_age_us); } pub fn flow_count(&self) -> usize { @@ -357,4 +366,4 @@ fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16) } None -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/inference.rs b/mantis/src/detection/ml/inference.rs index 227e731..04520fd 100644 --- a/mantis/src/detection/ml/inference.rs +++ b/mantis/src/detection/ml/inference.rs @@ -1,5 +1,5 @@ -use std::collections::VecDeque; use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use std::time::Instant; @@ -90,10 +90,7 @@ impl Inference { } let t2 = Instant::now(); - let ae_input = Array3::from_shape_fn( - (1, window_size, feat_len), - |(_, t, f)| sequence[t][f], - ); + let ae_input = Array3::from_shape_fn((1, window_size, feat_len), |(_, t, f)| sequence[t][f]); let t3 = Instant::now(); let ae_score = match self.run_autoencoder(&ae_input) { @@ -114,14 +111,19 @@ impl Inference { )); let pad = window_size - buf_len_snapshot; - let rows = sequence.iter().enumerate().map(|(t, row)| { - if t < pad { - format!(" t{t:02}: [pad]") - } else { - let vals = row.iter().map(|v| format!("{v:7.3}")).collect::>().join(" "); - format!(" t{t:02}: [{vals}]") - } - }).collect::>().join("\n"); + let rows = sequence + .iter() + .enumerate() + .map(|(t, row)| { + if t < pad { + format!(" t{t:02}: [pad]") + } else { + let vals = row.iter().map(|v| format!("{v:7.3}")).collect::>().join(" "); + format!(" t{t:02}: [{vals}]") + } + }) + .collect::>() + .join("\n"); log!(MLLog::WindowDebug( flow.flow_key.src_ip.clone(), pad, @@ -180,4 +182,4 @@ impl Inference { Ok(mse) } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/mod.rs b/mantis/src/detection/ml/mod.rs index b3ddcd4..dd7d715 100644 --- a/mantis/src/detection/ml/mod.rs +++ b/mantis/src/detection/ml/mod.rs @@ -1,8 +1,8 @@ -pub mod model_loader; -pub mod config_loader; -pub mod flow_tracker; -pub mod feature_extractor; -pub mod inference; -pub mod engine; pub mod aggregator; +pub mod config_loader; +pub mod engine; +pub mod feature_extractor; +pub mod flow_tracker; +pub mod inference; +pub mod model_loader; pub mod traffic_logger; diff --git a/mantis/src/detection/ml/model_loader.rs b/mantis/src/detection/ml/model_loader.rs index 87d57ea..99f2e31 100644 --- a/mantis/src/detection/ml/model_loader.rs +++ b/mantis/src/detection/ml/model_loader.rs @@ -18,7 +18,8 @@ impl MLModels { { log!(MLLog::BackendNativeOrt); ort::init_from(PathBuf::from(env!("ONNXRUNTIME_PATH"))) - .map_err(|_| MLError::InitializeFailed)?.commit(); + .map_err(|_| MLError::InitializeFailed)? + .commit(); } #[cfg(feature = "tract-backend")] { @@ -35,9 +36,13 @@ impl MLModels { let model_path = PathBuf::from(env!("ARTIFACTCS_PATH")).join(model_name); Session::builder() - .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })? + .map_err(|_| MLError::ModelLoadFailed { + path: model_path.clone(), + })? .with_optimization_level(GraphOptimizationLevel::All) - .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })? + .map_err(|_| MLError::ModelLoadFailed { + path: model_path.clone(), + })? .commit_from_file(&model_path) .map_err(|_| MLError::ModelLoadFailed { path: model_path }) } @@ -55,4 +60,4 @@ impl MLModels { _ => "unknown model".to_string(), } } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/traffic_logger.rs b/mantis/src/detection/ml/traffic_logger.rs index e2e4954..4013cbd 100644 --- a/mantis/src/detection/ml/traffic_logger.rs +++ b/mantis/src/detection/ml/traffic_logger.rs @@ -1,11 +1,11 @@ use std::fs::{self, OpenOptions}; use std::io::{BufWriter, Write}; -use std::path::{PathBuf, Path}; +use std::path::{Path, PathBuf}; use std::thread; use std::time::Duration; use chrono::Local; -use crossbeam::channel::{bounded, Sender, TrySendError, RecvTimeoutError}; +use crossbeam::channel::{RecvTimeoutError, Sender, TrySendError, bounded}; use macros::log; use crate::model::error::ml::MLError; diff --git a/mantis/src/detection/mod.rs b/mantis/src/detection/mod.rs index 12c65bf..ff61bce 100644 --- a/mantis/src/detection/mod.rs +++ b/mantis/src/detection/mod.rs @@ -1,3 +1,3 @@ pub mod fusion; pub mod ml; -pub mod suricata; \ No newline at end of file +pub mod suricata; diff --git a/mantis/src/detection/suricata/engine.rs b/mantis/src/detection/suricata/engine.rs index e704022..3d63bae 100644 --- a/mantis/src/detection/suricata/engine.rs +++ b/mantis/src/detection/suricata/engine.rs @@ -6,16 +6,15 @@ use std::sync::Arc; use std::thread; use std::time::Duration; -use crossbeam::channel::{bounded, Sender}; +use crossbeam::channel::{Sender, bounded}; use macros::log; +use super::output; use crate::detection::fusion::FusionEngine; use crate::model::config::SuricataConfig; use crate::model::error::suricata::SuricataError; use crate::model::log::suricata::SuricataLog; -use super::output; - const MIRROR_IFACE: &str = "mantis-mirror"; const MIRROR_PEER: &str = "mantis-peer"; const SURICATA_LOG: &str = "/tmp/suricata.log"; @@ -42,10 +41,12 @@ impl SuricataEngine { output::start_eve_reader(path, fusion); } - let rule_path_str = rule_path.to_str() - .ok_or_else(|| SuricataError::InvalidPath { path: rule_path.display().to_string() })?; - let eve_socket_str = eve_socket.to_str() - .ok_or_else(|| SuricataError::InvalidPath { path: eve_socket.display().to_string() })?; + let rule_path_str = rule_path.to_str().ok_or_else(|| SuricataError::InvalidPath { + path: rule_path.display().to_string(), + })?; + let eve_socket_str = eve_socket.to_str().ok_or_else(|| SuricataError::InvalidPath { + path: eve_socket.display().to_string(), + })?; let suppress = Self::generate_suppress(&config.suppress); let suppress_fd = Self::yaml_to_memfd(&suppress)?; @@ -73,7 +74,9 @@ impl SuricataEngine { .spawn(|| { let log_path = Path::new(SURICATA_LOG); for _ in 0..100 { - if log_path.exists() { break; } + if log_path.exists() { + break; + } thread::sleep(Duration::from_millis(100)); } let file = match std::fs::File::open(log_path) { @@ -96,17 +99,20 @@ impl SuricataEngine { if inotify_fd >= 0 { // Block until Suricata writes more data. let mut buf = [0u8; 64]; - let n = unsafe { - libc::read(inotify_fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) - }; - if n <= 0 { break; } + let n = + unsafe { libc::read(inotify_fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) }; + if n <= 0 { + break; + } } else { thread::sleep(Duration::from_millis(50)); } } Ok(_) => { let trimmed = line.trim_end().to_string(); - if trimmed.is_empty() { continue; } + if trimmed.is_empty() { + continue; + } let lower = trimmed.to_ascii_lowercase(); if lower.starts_with("error") || lower.starts_with("critical") { log!(SuricataLog::ProcessError { line: trimmed }); @@ -144,7 +150,9 @@ impl SuricataEngine { sll.sll_protocol = (libc::ETH_P_ALL as u16).to_be(); sll.sll_ifindex = ifindex as i32; - log!(SuricataLog::MirrorReady { iface: MIRROR_IFACE.into() }); + log!(SuricataLog::MirrorReady { + iface: MIRROR_IFACE.into() + }); while let Ok(data) = rx.recv() { unsafe { @@ -166,7 +174,10 @@ impl SuricataEngine { log!(SuricataLog::Initialized); - Ok(Arc::new(Self { tx, child: std::sync::Mutex::new(child) })) + Ok(Arc::new(Self { + tx, + child: std::sync::Mutex::new(child), + })) } /* Non-blocking: drops silently when the channel is full under load. */ @@ -187,25 +198,31 @@ impl SuricataEngine { fn generate_yaml(config: &SuricataConfig, rule_path: &str, eve_socket: &str, suppress_path: &str) -> String { let threading = match (config.worker_cpu_set, config.management_cpu) { (None, None) => r#"threading: - set-cpu-affinity: no"#.to_string(), + set-cpu-affinity: no"# + .to_string(), (worker, mgmt) => { let mgmt_cpu = mgmt.unwrap_or(0); let worker_block = match worker { - Some([start, end]) => format!(r#" - worker-cpu-set: + Some([start, end]) => format!( + r#" - worker-cpu-set: cpu: [ "{start}-{end}" ] - mode: "balanced""#), + mode: "balanced""# + ), None => String::new(), }; - format!(r#"threading: + format!( + r#"threading: set-cpu-affinity: yes cpu-affinity: - management-cpu-set: cpu: [ {mgmt_cpu} ] -{worker_block}"#) +{worker_block}"# + ) } }; - format!(r#"%YAML 1.1 + format!( + r#"%YAML 1.1 --- vars: address-groups: @@ -298,23 +315,21 @@ legacy: host-mode: sniffer-only "#, - home_net = config.home_net, - rule_path = rule_path, - eve_socket = eve_socket, - suppress_path = suppress_path, - log_path = SURICATA_LOG, - iface = MIRROR_PEER, - threads = config.af_packet_threads, - ring_size = config.af_packet_ring_size, - block_size = config.af_packet_block_size, - threading = threading, + home_net = config.home_net, + rule_path = rule_path, + eve_socket = eve_socket, + suppress_path = suppress_path, + log_path = SURICATA_LOG, + iface = MIRROR_PEER, + threads = config.af_packet_threads, + ring_size = config.af_packet_ring_size, + block_size = config.af_packet_block_size, + threading = threading, ) } fn yaml_to_memfd(yaml: &str) -> Result { - let fd = unsafe { - libc::memfd_create(b"suricata-config\0".as_ptr() as *const libc::c_char, 0) - }; + let fd = unsafe { libc::memfd_create(b"suricata-config\0".as_ptr() as *const libc::c_char, 0) }; if fd < 0 { let errno = unsafe { *libc::__errno_location() }; return Err(SuricataError::MirrorSetupFailed { @@ -322,9 +337,7 @@ host-mode: sniffer-only }); } let bytes = yaml.as_bytes(); - let written = unsafe { - libc::write(fd, bytes.as_ptr() as *const libc::c_void, bytes.len()) - }; + let written = unsafe { libc::write(fd, bytes.as_ptr() as *const libc::c_void, bytes.len()) }; if written < 0 { unsafe { libc::close(fd) }; let errno = unsafe { *libc::__errno_location() }; @@ -351,7 +364,10 @@ host-mode: sniffer-only .map_err(|e| SuricataError::MirrorSetupFailed { reason: e.to_string() })?; } - log!(SuricataLog::VethCreated { iface: MIRROR_IFACE.into(), peer: MIRROR_PEER.into() }); + log!(SuricataLog::VethCreated { + iface: MIRROR_IFACE.into(), + peer: MIRROR_PEER.into() + }); Ok(()) } @@ -366,13 +382,7 @@ host-mode: sniffer-only } fn open_raw_socket() -> Result { - let fd = unsafe { - libc::socket( - libc::AF_PACKET, - libc::SOCK_RAW, - (libc::ETH_P_ALL as u16).to_be() as i32, - ) - }; + let fd = unsafe { libc::socket(libc::AF_PACKET, libc::SOCK_RAW, (libc::ETH_P_ALL as u16).to_be() as i32) }; if fd < 0 { let errno = unsafe { *libc::__errno_location() }; return Err(SuricataError::MirrorSetupFailed { @@ -390,4 +400,4 @@ impl Drop for SuricataEngine { } let _ = Command::new("ip").args(["link", "del", MIRROR_IFACE]).output(); } -} \ No newline at end of file +} diff --git a/mantis/src/detection/suricata/mod.rs b/mantis/src/detection/suricata/mod.rs index e1095f0..82fbe20 100644 --- a/mantis/src/detection/suricata/mod.rs +++ b/mantis/src/detection/suricata/mod.rs @@ -1,3 +1,3 @@ -mod output; pub mod engine; +mod output; pub use engine::SuricataEngine; diff --git a/mantis/src/detection/suricata/output.rs b/mantis/src/detection/suricata/output.rs index a3eb419..6024488 100644 --- a/mantis/src/detection/suricata/output.rs +++ b/mantis/src/detection/suricata/output.rs @@ -78,9 +78,7 @@ fn handle_eve_stream(stream: std::os::unix::net::UnixStream, fusion: Arc Result<(), Error> { let mut system = System::new().await?; system.run().await?; system.terminate().await?; - + drop(system); std::process::exit(0); -} \ No newline at end of file +} diff --git a/mantis/src/model/config.rs b/mantis/src/model/config.rs index 88422c8..7f4015a 100644 --- a/mantis/src/model/config.rs +++ b/mantis/src/model/config.rs @@ -21,9 +21,15 @@ pub struct SuricataConfig { pub suppress: Vec, } -fn default_af_threads() -> String { "auto".to_string() } -fn default_af_ring_size() -> u32 { 2048 } -fn default_af_block_size() -> u32 { 131072 } +fn default_af_threads() -> String { + "auto".to_string() +} +fn default_af_ring_size() -> u32 { + 2048 +} +fn default_af_block_size() -> u32 { + 131072 +} #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Config { @@ -75,8 +81,12 @@ pub struct AuthConfig { pub default_admin_password: String, } -fn default_token_ttl() -> u64 { 86400 } -fn default_admin_password() -> String { "admin".to_string() } +fn default_token_ttl() -> u64 { + 86400 +} +fn default_admin_password() -> String { + "admin".to_string() +} fn default_fusion_mode() -> String { "or".to_string() @@ -88,4 +98,4 @@ fn default_fusion_window_secs() -> u64 { fn default_ae_threshold_method() -> String { "95".to_string() -} \ No newline at end of file +} diff --git a/mantis/src/model/error/ebpf.rs b/mantis/src/model/error/ebpf.rs index c59a1f7..dfb2387 100644 --- a/mantis/src/model/error/ebpf.rs +++ b/mantis/src/model/error/ebpf.rs @@ -70,4 +70,3 @@ traceable! { FlowMapKeyMissing { direction: String, flow_direction: String, time_type: String } => tracing::Level::ERROR, } } - diff --git a/mantis/src/model/error/misc.rs b/mantis/src/model/error/misc.rs index cf7db5a..886f716 100644 --- a/mantis/src/model/error/misc.rs +++ b/mantis/src/model/error/misc.rs @@ -26,4 +26,4 @@ traceable! { #[error("Failed to create traffic log file '{path}': {reason}")] TrafficLogCreateError { path: String, reason: String } => tracing::Level::ERROR, } -} \ No newline at end of file +} diff --git a/mantis/src/model/error/ml.rs b/mantis/src/model/error/ml.rs index a22bb4b..752d73a 100644 --- a/mantis/src/model/error/ml.rs +++ b/mantis/src/model/error/ml.rs @@ -48,4 +48,4 @@ traceable! { #[error("Failed to write CSV row: {reason}")] TrafficLogWriteFailed { reason: String } => tracing::Level::WARN, } -} \ No newline at end of file +} diff --git a/mantis/src/model/error/mod.rs b/mantis/src/model/error/mod.rs index 255f75c..be1b678 100644 --- a/mantis/src/model/error/mod.rs +++ b/mantis/src/model/error/mod.rs @@ -4,8 +4,8 @@ pub mod http; pub mod io; pub mod misc; pub mod ml; -pub mod system; pub mod suricata; +pub mod system; use serde::{Deserialize, Serialize}; @@ -84,4 +84,4 @@ impl From for Error { fn from(error: SuricataError) -> Self { Self::Suricata(error) } -} \ No newline at end of file +} diff --git a/mantis/src/model/geo_stats.rs b/mantis/src/model/geo_stats.rs index 4bf850e..b92c864 100644 --- a/mantis/src/model/geo_stats.rs +++ b/mantis/src/model/geo_stats.rs @@ -1,5 +1,5 @@ -use serde::{Deserialize, Serialize}; use common::model::flow_stats::FlowStats; +use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize)] pub struct FlowStatsWithGeo { diff --git a/mantis/src/model/ip_address.rs b/mantis/src/model/ip_address.rs index 488a082..06b4999 100644 --- a/mantis/src/model/ip_address.rs +++ b/mantis/src/model/ip_address.rs @@ -37,17 +37,11 @@ impl NativeConvert for AddrPortV4 { type Native = SocketAddrV4; fn into_native(self) -> Self::Native { - SocketAddrV4::new( - Ipv4Addr::from(u32::from_be(self.ip())), - self.port() - ) + SocketAddrV4::new(Ipv4Addr::from(u32::from_be(self.ip())), self.port()) } fn from_native(native: Self::Native) -> Self { - AddrPortV4::new( - native.ip().to_bits().to_be(), - native.port() - ) + AddrPortV4::new(native.ip().to_bits().to_be(), native.port()) } } @@ -55,18 +49,10 @@ impl NativeConvert for AddrPortV6 { type Native = SocketAddrV6; fn into_native(self) -> Self::Native { - SocketAddrV6::new( - Ipv6Addr::from(u128::from_be(self.ip())), - self.port(), - 0, - 0 - ) + SocketAddrV6::new(Ipv6Addr::from(u128::from_be(self.ip())), self.port(), 0, 0) } fn from_native(native: Self::Native) -> Self { - AddrPortV6::new( - native.ip().to_bits().to_be(), - native.port() - ) + AddrPortV6::new(native.ip().to_bits().to_be(), native.port()) } -} \ No newline at end of file +} diff --git a/mantis/src/model/log/ebpf.rs b/mantis/src/model/log/ebpf.rs index 3f6c6dc..3cd72aa 100644 --- a/mantis/src/model/log/ebpf.rs +++ b/mantis/src/model/log/ebpf.rs @@ -26,7 +26,7 @@ loggable! { #[error("No frames available for TX")] NoFramesAvailable => tracing::Level::WARN, - + #[error("TX wakeup failed: {error}")] TXWakeupFailed { error: String } => tracing::Level::WARN, @@ -51,4 +51,4 @@ loggable! { #[error("Fill queue incomplete: produced {produced}, expected {expected}")] FillQueueIncomplete { produced: usize, expected: usize } => tracing::Level::WARN, } -} \ No newline at end of file +} diff --git a/mantis/src/model/log/ml.rs b/mantis/src/model/log/ml.rs index a113b52..39982d6 100644 --- a/mantis/src/model/log/ml.rs +++ b/mantis/src/model/log/ml.rs @@ -79,4 +79,4 @@ loggable! { WindowDebug { src: String, pad: usize, window_size: usize, ae_score: f32, rows: String } => tracing::Level::DEBUG, } -} \ No newline at end of file +} diff --git a/mantis/src/model/log/mod.rs b/mantis/src/model/log/mod.rs index ce8339f..32a26ac 100644 --- a/mantis/src/model/log/mod.rs +++ b/mantis/src/model/log/mod.rs @@ -1,8 +1,8 @@ pub mod auth; pub mod ebpf; -pub mod http; -pub mod ml; -pub mod system; -pub mod misc; pub mod health; +pub mod http; +pub mod misc; +pub mod ml; pub mod suricata; +pub mod system; diff --git a/mantis/src/model/log/suricata.rs b/mantis/src/model/log/suricata.rs index ebbd808..35e1355 100644 --- a/mantis/src/model/log/suricata.rs +++ b/mantis/src/model/log/suricata.rs @@ -33,4 +33,4 @@ loggable! { #[error("[suricata] {line}")] ProcessError { line: String } => tracing::Level::ERROR, } -} \ No newline at end of file +} diff --git a/mantis/src/model/log/system.rs b/mantis/src/model/log/system.rs index 22588c7..c4bce53 100644 --- a/mantis/src/model/log/system.rs +++ b/mantis/src/model/log/system.rs @@ -28,4 +28,4 @@ loggable! { TrafficLoggingEnabled { path: String } => tracing::Level::INFO, } -} \ No newline at end of file +} diff --git a/mantis/src/model/ml_detection.rs b/mantis/src/model/ml_detection.rs index 74d224f..c383ffb 100644 --- a/mantis/src/model/ml_detection.rs +++ b/mantis/src/model/ml_detection.rs @@ -236,4 +236,4 @@ impl UnifiedAlert { rule_msg: Some(m.msg.clone()), } } -} \ No newline at end of file +} diff --git a/mantis/src/model/mod.rs b/mantis/src/model/mod.rs index 277a82e..566f6f0 100644 --- a/mantis/src/model/mod.rs +++ b/mantis/src/model/mod.rs @@ -2,10 +2,10 @@ pub mod config; pub mod direction; pub mod error; pub mod geo_stats; +pub mod health; pub mod ip_address; pub mod list_type; pub mod log; -pub mod time_type; pub mod ml_detection; -pub mod health; pub mod rule_detection; +pub mod time_type; diff --git a/mantis/src/utils/ip_address.rs b/mantis/src/utils/ip_address.rs index 2203b23..e8e326b 100644 --- a/mantis/src/utils/ip_address.rs +++ b/mantis/src/utils/ip_address.rs @@ -13,12 +13,7 @@ pub fn convert_ports_to_vec(ports: [u16; MAX_RULES_PORT]) -> Vec { pub fn is_private_ip(ip: &IpAddr) -> bool { match ip { - IpAddr::V4(v4) => { - v4.is_private() - || v4.is_loopback() - || v4.is_link_local() - || v4.is_broadcast() - } + IpAddr::V4(v4) => v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_broadcast(), IpAddr::V6(v6) => { v6.is_loopback() || v6.is_unique_local() // fc00::/7 @@ -26,4 +21,4 @@ pub fn is_private_ip(ip: &IpAddr) -> bool { || v6.is_multicast() } } -} \ No newline at end of file +} diff --git a/mantis/src/utils/logging.rs b/mantis/src/utils/logging.rs index b88090c..18b0904 100644 --- a/mantis/src/utils/logging.rs +++ b/mantis/src/utils/logging.rs @@ -1,20 +1,20 @@ use std::fs; + use tracing::Level; use tracing_appender::rolling::{RollingFileAppender, Rotation}; use tracing_subscriber::filter::EnvFilter; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; -use crate::model::error::io::IOError; use crate::model::error::Error; +use crate::model::error::io::IOError; pub struct Logging; impl Logging { pub fn initialize() -> Result<(), Error> { let log_directory = "logs"; - fs::create_dir_all(log_directory) - .map_err(|err| IOError::CreateDirectoryFailed(log_directory, err))?; + fs::create_dir_all(log_directory).map_err(|err| IOError::CreateDirectoryFailed(log_directory, err))?; let file_appender = RollingFileAppender::new(Rotation::DAILY, log_directory, "Mantis"); diff --git a/mantis/src/utils/mod.rs b/mantis/src/utils/mod.rs index 50d5f76..b3ee894 100644 --- a/mantis/src/utils/mod.rs +++ b/mantis/src/utils/mod.rs @@ -1,7 +1,7 @@ -pub mod logging; -pub mod static_files; pub mod boot_time; -pub mod packet_parser; pub mod cpu_affinity; +pub mod logging; +pub mod packet_parser; +pub mod static_files; -pub mod ip_address; \ No newline at end of file +pub mod ip_address; diff --git a/mantis/src/utils/packet_parser.rs b/mantis/src/utils/packet_parser.rs index a24fdb0..be4f2a9 100644 --- a/mantis/src/utils/packet_parser.rs +++ b/mantis/src/utils/packet_parser.rs @@ -184,4 +184,4 @@ pub fn format_ipv6(addr: u128) -> String { bytes[14], bytes[15] ) -} \ No newline at end of file +} diff --git a/mantis/src/web/api/auth.rs b/mantis/src/web/api/auth.rs index 6e95b05..328d055 100644 --- a/mantis/src/web/api/auth.rs +++ b/mantis/src/web/api/auth.rs @@ -1,10 +1,13 @@ -use axum::{Router, routing::{get, post}}; +use axum::Json; use axum::extract::State; use axum::http::StatusCode; use axum::response::IntoResponse; -use axum::Json; +use axum::{ + Router, + routing::{get, post}, +}; use chrono::Utc; -use jsonwebtoken::{encode, EncodingKey, Header}; +use jsonwebtoken::{EncodingKey, Header, encode}; use macros::log; use serde::{Deserialize, Serialize}; @@ -38,10 +41,7 @@ struct MeResponse { role: String, } -async fn login( - State(state): State, - Json(body): Json, -) -> impl IntoResponse { +async fn login(State(state): State, Json(body): Json) -> impl IntoResponse { let auth_cfg = match state.app_config.auth.as_ref() { Some(c) => c, None => return (StatusCode::NOT_IMPLEMENTED, "Auth not configured").into_response(), @@ -54,14 +54,18 @@ async fn login( let account = match db.find_account_by_username(&body.username) { Ok(Some(a)) => a, Ok(None) => { - log!(AuthLog::LoginFailed { username: body.username.clone() }); + log!(AuthLog::LoginFailed { + username: body.username.clone() + }); return (StatusCode::UNAUTHORIZED, "Invalid credentials").into_response(); } Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response(), }; if !verify_password(&body.password, &account.password_hash) { - log!(AuthLog::LoginFailed { username: body.username.clone() }); + log!(AuthLog::LoginFailed { + username: body.username.clone() + }); return (StatusCode::UNAUTHORIZED, "Invalid credentials").into_response(); } @@ -79,7 +83,9 @@ async fn login( &EncodingKey::from_secret(auth_cfg.jwt_secret.as_bytes()), ) { Ok(token) => { - log!(AuthLog::LoginSuccess { username: account.username }); + log!(AuthLog::LoginSuccess { + username: account.username + }); Json(LoginResponse { token }).into_response() } Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, "Token generation failed").into_response(), diff --git a/mantis/src/web/api/control/access_control.rs b/mantis/src/web/api/control/access_control.rs index a84ba7b..600bb33 100644 --- a/mantis/src/web/api/control/access_control.rs +++ b/mantis/src/web/api/control/access_control.rs @@ -63,7 +63,11 @@ async fn remove_ipv4_list( State(state): State, Json(address): Json, ) -> impl IntoResponse { - match state.access_control.remove_ipv4_list(direction, list_type, address).await { + match state + .access_control + .remove_ipv4_list(direction, list_type, address) + .await + { Ok(_) => StatusCode::OK.into_response(), Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), } @@ -74,8 +78,12 @@ async fn remove_ipv6_list( State(state): State, Json(address): Json, ) -> impl IntoResponse { - match state.access_control.remove_ipv6_list(direction, list_type, address).await { + match state + .access_control + .remove_ipv6_list(direction, list_type, address) + .await + { Ok(_) => StatusCode::OK.into_response(), Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), } -} \ No newline at end of file +} diff --git a/mantis/src/web/api/control/service.rs b/mantis/src/web/api/control/service.rs index a28ef28..fb5a73d 100644 --- a/mantis/src/web/api/control/service.rs +++ b/mantis/src/web/api/control/service.rs @@ -11,17 +11,57 @@ use crate::core::app_state::AppState; pub fn router() -> Router { Router::new() - .route("/ipv4/http_service", get(get_ipv4_http_service).put(add_ipv4_http_service).delete(remove_ipv4_http_service)) - .route("/ipv6/http_service", get(get_ipv6_http_service).put(add_ipv6_http_service).delete(remove_ipv6_http_service)) + .route( + "/ipv4/http_service", + get(get_ipv4_http_service) + .put(add_ipv4_http_service) + .delete(remove_ipv4_http_service), + ) + .route( + "/ipv6/http_service", + get(get_ipv6_http_service) + .put(add_ipv6_http_service) + .delete(remove_ipv6_http_service), + ) .route("/ssh_white_list", get(is_ssh_white_list_enable)) .route("/ssh_white_list/enable", post(enable_ssh_white_list)) .route("/ssh_white_list/disable", post(disable_ssh_white_list)) - .route("/ipv4/ssh_service", get(get_ipv4_ssh_service).put(add_ipv4_ssh_service).delete(remove_ipv4_ssh_service)) - .route("/ipv6/ssh_service", get(get_ipv6_ssh_service).put(add_ipv6_ssh_service).delete(remove_ipv6_ssh_service)) - .route("/ipv4/ssh_white_list", get(get_ipv4_ssh_white_list).put(add_ipv4_ssh_white_list).delete(remove_ipv4_ssh_white_list)) - .route("/ipv6/ssh_white_list", get(get_ipv6_ssh_white_list).put(add_ipv6_ssh_white_list).delete(remove_ipv6_ssh_white_list)) - .route("/ipv4/ssh_black_list", get(get_ipv4_ssh_black_list).put(add_ipv4_ssh_black_list).delete(remove_ipv4_ssh_black_list)) - .route("/ipv6/ssh_black_list", get(get_ipv6_ssh_black_list).put(add_ipv6_ssh_black_list).delete(remove_ipv6_ssh_black_list)) + .route( + "/ipv4/ssh_service", + get(get_ipv4_ssh_service) + .put(add_ipv4_ssh_service) + .delete(remove_ipv4_ssh_service), + ) + .route( + "/ipv6/ssh_service", + get(get_ipv6_ssh_service) + .put(add_ipv6_ssh_service) + .delete(remove_ipv6_ssh_service), + ) + .route( + "/ipv4/ssh_white_list", + get(get_ipv4_ssh_white_list) + .put(add_ipv4_ssh_white_list) + .delete(remove_ipv4_ssh_white_list), + ) + .route( + "/ipv6/ssh_white_list", + get(get_ipv6_ssh_white_list) + .put(add_ipv6_ssh_white_list) + .delete(remove_ipv6_ssh_white_list), + ) + .route( + "/ipv4/ssh_black_list", + get(get_ipv4_ssh_black_list) + .put(add_ipv4_ssh_black_list) + .delete(remove_ipv4_ssh_black_list), + ) + .route( + "/ipv6/ssh_black_list", + get(get_ipv6_ssh_black_list) + .put(add_ipv6_ssh_black_list) + .delete(remove_ipv6_ssh_black_list), + ) } async fn get_ipv4_http_service(State(state): State) -> impl IntoResponse { @@ -32,7 +72,10 @@ async fn get_ipv6_http_service(State(state): State) -> impl IntoRespon Json(state.service.get_ipv6_http_service().await) } -async fn add_ipv4_http_service(State(state): State, Json(payload): Json<(SocketAddrV4, Vec)>) -> impl IntoResponse { +async fn add_ipv4_http_service( + State(state): State, + Json(payload): Json<(SocketAddrV4, Vec)>, +) -> impl IntoResponse { let (addr, methods) = payload; match state.service.add_ipv4_http_service(addr, methods).await { Ok(_) => StatusCode::OK.into_response(), @@ -40,7 +83,10 @@ async fn add_ipv4_http_service(State(state): State, Json(payload): Jso } } -async fn add_ipv6_http_service(State(state): State, Json(payload): Json<(SocketAddrV6, Vec)>) -> impl IntoResponse { +async fn add_ipv6_http_service( + State(state): State, + Json(payload): Json<(SocketAddrV6, Vec)>, +) -> impl IntoResponse { let (addr, methods) = payload; match state.service.add_ipv6_http_service(addr, methods).await { Ok(_) => StatusCode::OK.into_response(), @@ -48,7 +94,10 @@ async fn add_ipv6_http_service(State(state): State, Json(payload): Jso } } -async fn remove_ipv4_http_service(State(state): State, Json(payload): Json<(SocketAddrV4, Vec)>) -> impl IntoResponse { +async fn remove_ipv4_http_service( + State(state): State, + Json(payload): Json<(SocketAddrV4, Vec)>, +) -> impl IntoResponse { let (addr, methods) = payload; match state.service.remove_ipv4_http_service(addr, methods).await { Ok(_) => StatusCode::OK.into_response(), @@ -56,7 +105,10 @@ async fn remove_ipv4_http_service(State(state): State, Json(payload): } } -async fn remove_ipv6_http_service(State(state): State, Json(payload): Json<(SocketAddrV6, Vec)>) -> impl IntoResponse { +async fn remove_ipv6_http_service( + State(state): State, + Json(payload): Json<(SocketAddrV6, Vec)>, +) -> impl IntoResponse { let (addr, methods) = payload; match state.service.remove_ipv6_http_service(addr, methods).await { Ok(_) => StatusCode::OK.into_response(), diff --git a/mantis/src/web/api/control/statistics.rs b/mantis/src/web/api/control/statistics.rs index 2b1ad15..38198ec 100644 --- a/mantis/src/web/api/control/statistics.rs +++ b/mantis/src/web/api/control/statistics.rs @@ -1,5 +1,5 @@ -use axum::extract::{Path, State}; use axum::extract::ws::WebSocketUpgrade; +use axum::extract::{Path, State}; use axum::http::StatusCode; use axum::response::IntoResponse; use axum::routing::get; @@ -15,15 +15,25 @@ pub fn router() -> Router { Router::new() .route("/get/ipv4/{direction}/{flow_direction}/{time_type}", get(get_ipv4_flow)) .route("/get/ipv6/{direction}/{flow_direction}/{time_type}", get(get_ipv6_flow)) - .route("/websocket/ipv4/{direction}/{flow_direction}/{time_type}", get(websocket_ipv4)) - .route("/websocket/ipv6/{direction}/{flow_direction}/{time_type}", get(websocket_ipv6)) + .route( + "/websocket/ipv4/{direction}/{flow_direction}/{time_type}", + get(websocket_ipv4), + ) + .route( + "/websocket/ipv6/{direction}/{flow_direction}/{time_type}", + get(websocket_ipv6), + ) } async fn get_ipv4_flow( Path((direction, flow_direction, time_type)): Path<(Direction, FlowDirection, TimeType)>, State(state): State, ) -> impl IntoResponse { - match state.statistics.get_ipv4_flow_data(direction, flow_direction, time_type).await { + match state + .statistics + .get_ipv4_flow_data(direction, flow_direction, time_type) + .await + { Ok(data) => Json(data).into_response(), Err(e) => { log!(e); @@ -36,7 +46,11 @@ async fn get_ipv6_flow( Path((direction, flow_direction, time_type)): Path<(Direction, FlowDirection, TimeType)>, State(state): State, ) -> impl IntoResponse { - match state.statistics.get_ipv6_flow_data(direction, flow_direction, time_type).await { + match state + .statistics + .get_ipv6_flow_data(direction, flow_direction, time_type) + .await + { Ok(data) => Json(data).into_response(), Err(e) => { log!(e); @@ -50,9 +64,7 @@ async fn websocket_ipv4( State(state): State, ws: WebSocketUpgrade, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| { - flow_websocket::handle_ipv4_flow(socket, state, direction, flow_direction, time_type) - }) + ws.on_upgrade(move |socket| flow_websocket::handle_ipv4_flow(socket, state, direction, flow_direction, time_type)) } async fn websocket_ipv6( @@ -60,7 +72,5 @@ async fn websocket_ipv6( State(state): State, ws: WebSocketUpgrade, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| { - flow_websocket::handle_ipv6_flow(socket, state, direction, flow_direction, time_type) - }) -} \ No newline at end of file + ws.on_upgrade(move |socket| flow_websocket::handle_ipv6_flow(socket, state, direction, flow_direction, time_type)) +} diff --git a/mantis/src/web/api/default.rs b/mantis/src/web/api/default.rs index 3d8a5cf..b0720e0 100644 --- a/mantis/src/web/api/default.rs +++ b/mantis/src/web/api/default.rs @@ -13,8 +13,7 @@ pub async fn default_route(uri: Uri) -> Response { }; if let Some(content) = StaticFiles::get(&file_system_path) { - let mime_type = mime_guess::from_path(&file_system_path) - .first_or_octet_stream(); + let mime_type = mime_guess::from_path(&file_system_path).first_or_octet_stream(); return ( [("content-type", mime_type.as_ref().to_string())], content.data.into_owned(), @@ -24,20 +23,12 @@ pub async fn default_route(uri: Uri) -> Response { let html_path = format!("{}.html", file_system_path); if let Some(content) = StaticFiles::get(&html_path) { - return ( - [("content-type", "text/html")], - content.data.into_owned(), - ) - .into_response(); + return ([("content-type", "text/html")], content.data.into_owned()).into_response(); } let index_path = format!("{}/index.html", file_system_path); if let Some(content) = StaticFiles::get(&index_path) { - return ( - [("content-type", "text/html")], - content.data.into_owned(), - ) - .into_response(); + return ([("content-type", "text/html")], content.data.into_owned()).into_response(); } match StaticFiles::get("web/404.html") { diff --git a/mantis/src/web/api/detection_alert.rs b/mantis/src/web/api/detection_alert.rs index 903959a..9f74f04 100644 --- a/mantis/src/web/api/detection_alert.rs +++ b/mantis/src/web/api/detection_alert.rs @@ -1,6 +1,6 @@ -use axum::{Router, routing::get}; use axum::extract::State; use axum::response::IntoResponse; +use axum::{Router, routing::get}; use crate::core::app_state::AppState; use crate::web::websocket::alert_websocket; @@ -9,10 +9,7 @@ pub fn router() -> Router { Router::new().route("/websocket/alert", get(websocket_alert)) } -async fn websocket_alert( - ws: axum::extract::ws::WebSocketUpgrade, - State(state): State, -) -> impl IntoResponse { +async fn websocket_alert(ws: axum::extract::ws::WebSocketUpgrade, State(state): State) -> impl IntoResponse { let rx = state.detection_alert.subscribe(); ws.on_upgrade(|socket| alert_websocket::handle_alert(socket, rx)) } diff --git a/mantis/src/web/api/health.rs b/mantis/src/web/api/health.rs index 7134a42..7872464 100644 --- a/mantis/src/web/api/health.rs +++ b/mantis/src/web/api/health.rs @@ -1,7 +1,7 @@ -use axum::{Router, routing::get}; +use axum::Json; use axum::extract::State; use axum::response::IntoResponse; -use axum::Json; +use axum::{Router, routing::get}; use crate::core::app_state::AppState; use crate::web::websocket::health_websocket; diff --git a/mantis/src/web/api/misc.rs b/mantis/src/web/api/misc.rs index 94f4a48..55f87cf 100644 --- a/mantis/src/web/api/misc.rs +++ b/mantis/src/web/api/misc.rs @@ -1,6 +1,6 @@ -use axum::{Router, routing::get}; -use axum::response::IntoResponse; use axum::Json; +use axum::response::IntoResponse; +use axum::{Router, routing::get}; use crate::core::app_state::AppState; use crate::utils::boot_time::boot_time; diff --git a/mantis/src/web/middleware/auth.rs b/mantis/src/web/middleware/auth.rs index d6bb827..39e1029 100644 --- a/mantis/src/web/middleware/auth.rs +++ b/mantis/src/web/middleware/auth.rs @@ -1,6 +1,6 @@ use axum::extract::FromRequestParts; -use axum::http::{request::Parts, StatusCode}; -use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation}; +use axum::http::{StatusCode, request::Parts}; +use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode}; use serde::{Deserialize, Serialize}; use crate::core::app_state::AppState; @@ -18,10 +18,7 @@ pub struct AuthenticatedUser(pub Claims); impl FromRequestParts for AuthenticatedUser { type Rejection = (StatusCode, &'static str); - async fn from_request_parts( - parts: &mut Parts, - state: &AppState, - ) -> Result { + async fn from_request_parts(parts: &mut Parts, state: &AppState) -> Result { let auth_cfg = state .app_config .auth diff --git a/mantis/src/web/websocket/alert_websocket.rs b/mantis/src/web/websocket/alert_websocket.rs index b81cfa8..074f0e4 100644 --- a/mantis/src/web/websocket/alert_websocket.rs +++ b/mantis/src/web/websocket/alert_websocket.rs @@ -3,10 +3,10 @@ use futures_util::{SinkExt, StreamExt}; use macros::log; use tokio::sync::broadcast; -use crate::model::ml_detection::UnifiedAlert; use crate::model::error::http::HttpError; use crate::model::error::misc::MiscError; use crate::model::log::http::HttpLog; +use crate::model::ml_detection::UnifiedAlert; pub async fn handle_alert(socket: WebSocket, mut broadcast_rx: broadcast::Receiver) { let (mut sender, mut receiver) = socket.split(); diff --git a/mantis/src/web/websocket/flow_websocket.rs b/mantis/src/web/websocket/flow_websocket.rs index 7ee255f..a29bf2e 100644 --- a/mantis/src/web/websocket/flow_websocket.rs +++ b/mantis/src/web/websocket/flow_websocket.rs @@ -1,7 +1,7 @@ use axum::extract::ws::{Message, WebSocket}; use futures_util::{SinkExt, StreamExt}; use macros::log; -use tokio::time::{interval, Duration}; +use tokio::time::{Duration, interval}; use crate::core::app_state::AppState; use crate::model::direction::{Direction, FlowDirection}; diff --git a/mantis/src/web/websocket/health_websocket.rs b/mantis/src/web/websocket/health_websocket.rs index aa73225..a93f474 100644 --- a/mantis/src/web/websocket/health_websocket.rs +++ b/mantis/src/web/websocket/health_websocket.rs @@ -3,9 +3,9 @@ use futures_util::{SinkExt, StreamExt}; use macros::log; use tokio::sync::broadcast; -use crate::model::health::SystemHealthMetrics; use crate::model::error::http::HttpError; use crate::model::error::misc::MiscError; +use crate::model::health::SystemHealthMetrics; use crate::model::log::http::HttpLog; pub async fn handle_health(socket: WebSocket, mut broadcast_rx: broadcast::Receiver) { diff --git a/mantis/src/web/websocket/mod.rs b/mantis/src/web/websocket/mod.rs index e8139e1..8ae1755 100644 --- a/mantis/src/web/websocket/mod.rs +++ b/mantis/src/web/websocket/mod.rs @@ -1,3 +1,3 @@ +pub mod alert_websocket; pub mod flow_websocket; pub mod health_websocket; -pub mod alert_websocket;