ParrotXray/Mantis
1
0
Fork 0
mirror of https://github.com/ParrotXray/Mantis.git synced 2026-08-24 18:50:28 +09:00
Code Issues Packages Projects Releases Wiki Activity

feat: adjust code with rustfmt

Browse Source
This commit is contained in:
ParrotXray 2026-05-23 03:27:21 +00:00
parent 1b6ac4bcee
commit 7d900c5dd8
59 changed files with 488 additions and 393 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
common/src/model/event.rs
View File

@ -109,7 +109,6 @@ impl Event {
}
}
#[repr(C, align(8))]
#[derive(Debug, Clone)]
pub struct IPv4Event {

4
macros/src/log.rs
View File

@ -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()
}

2
mantis/src/core/ebpf/access_control.rs
View File

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

10
mantis/src/core/ebpf/mod.rs
View File

@ -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<XskManager>,
@ -29,11 +29,7 @@ pub struct EbpfServices {
}
impl EbpfServices {
pub fn new(
app_config: Arc<AppConfig>,
ingress_ebpf: &mut Ebpf,
egress_ebpf: &mut Ebpf,
) -> Result<Self, Error> {
pub fn new(app_config: Arc<AppConfig>, ingress_ebpf: &mut Ebpf, egress_ebpf: &mut Ebpf) -> Result<Self, Error> {
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 {
}
}
}
}
}

2
mantis/src/core/ebpf/service.rs
View File

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

24
mantis/src/core/ebpf/statistics.rs
View File

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

14
mantis/src/core/ebpf/xsk_manager.rs
View File

@ -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<Vec<u8>>,
) -> Result<usize, EbpfError> {
fn process_rx_queue(&mut self, forward_tx: &Sender<Vec<u8>>) -> Result<usize, EbpfError> {
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)
}
}
}

2
mantis/src/core/infrastructure/app_config.rs
View File

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

16
mantis/src/core/infrastructure/app_db.rs
View File

@ -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) })
}

47
mantis/src/core/infrastructure/geoip.rs
View File

@ -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<P: AsRef<Path>>(
db_path: P,
cache_size: usize,
) -> Result<Self, MaxMindDbError> {
pub fn with_cache_size<P: AsRef<Path>>(db_path: P, cache_size: usize) -> Result<Self, MaxMindDbError> {
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<Vec<u8>>,
ip: IpAddr,
) -> Result<Option<GeoLocation>, MaxMindDbError> {
fn lookup_from_db_blocking(reader: &Reader<Vec<u8>>, ip: IpAddr) -> Result<Option<GeoLocation>, MaxMindDbError> {
let lookup_result = reader.lookup(ip)?;
let city_option: Option<geoip2::City> = 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())
}
}
}

58
mantis/src/core/infrastructure/health.rs
View File

@ -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<System>,
@ -30,7 +23,6 @@ pub struct SystemHealth {
egress_interface: String,
}
impl SystemHealth {
pub fn new(config: Arc<AppConfig>) -> Result<Self, Error> {
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
}
}
}

26
mantis/src/core/infrastructure/mod.rs
View File

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

12
mantis/src/core/system.rs
View File

@ -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<AppConfig>,
@ -243,4 +243,4 @@ impl System {
}
Ok(())
}
}
}

24
mantis/src/detection/fusion.rs
View File

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

22
mantis/src/detection/ml/aggregator.rs
View File

@ -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::<f32>() / detections.len() as f32;
let avg_score: f32 = detections.iter().map(|(_, s)| s).sum::<f32>() / 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()
}
}
}

20
mantis/src/detection/ml/config_loader.rs
View File

@ -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<Self, MLError> {
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()
}
}
}

37
mantis/src/detection/ml/engine.rs
View File

@ -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<String> = flows
.iter()
.map(|f| f.flow_key.src_ip.clone())
.collect();
let active_ips: std::collections::HashSet<String> =
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 }
}
}
}

31
mantis/src/detection/ml/feature_extractor.rs
View File

@ -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<f64> = 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<f64> = flow
.fwd_packets
.iter()
.map(|p| p.header_length as f64)
.collect();
let fwd_header_sizes: Vec<f64> = 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::<Vec<_>>());
@ -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()
}
}

49
mantis/src/detection/ml/flow_tracker.rs
View File

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

30
mantis/src/detection/ml/inference.rs
View File

@ -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::<Vec<_>>().join(" ");
format!(" t{t:02}: [{vals}]")
}
}).collect::<Vec<_>>().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::<Vec<_>>().join(" ");
format!(" t{t:02}: [{vals}]")
}
})
.collect::<Vec<_>>()
.join("\n");
log!(MLLog::WindowDebug(
flow.flow_key.src_ip.clone(),
pad,
@ -180,4 +182,4 @@ impl Inference {
Ok(mse)
}
}
}

12
mantis/src/detection/ml/mod.rs
View File

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

13
mantis/src/detection/ml/model_loader.rs
View File

@ -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(),
}
}
}
}

4
mantis/src/detection/ml/traffic_logger.rs
View File

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

2
mantis/src/detection/mod.rs
View File

@ -1,3 +1,3 @@
pub mod fusion;
pub mod ml;
pub mod suricata;
pub mod suricata;

102
mantis/src/detection/suricata/engine.rs
View File

@ -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<i32, SuricataError> {
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<i32, SuricataError> {
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();
}
}
}

2
mantis/src/detection/suricata/mod.rs
View File

@ -1,3 +1,3 @@
mod output;
pub mod engine;
mod output;
pub use engine::SuricataEngine;

6
mantis/src/detection/suricata/output.rs
View File

@ -78,9 +78,7 @@ fn handle_eve_stream(stream: std::os::unix::net::UnixStream, fusion: Arc<FusionE
if event.event_type != "alert" {
continue;
}
let (Some(alert), Some(src_ip), Some(dst_ip)) =
(event.alert, event.src_ip, event.dest_ip)
else {
let (Some(alert), Some(src_ip), Some(dst_ip)) = (event.alert, event.src_ip, event.dest_ip) else {
continue;
};
@ -103,4 +101,4 @@ fn handle_eve_stream(stream: std::os::unix::net::UnixStream, fusion: Arc<FusionE
fusion.record_rule(&m);
}
}
}

6
mantis/src/main.rs
View File

@ -1,8 +1,8 @@
mod core;
mod detection;
mod model;
mod utils;
mod web;
mod detection;
use crate::core::system::System;
use crate::model::error::Error;
@ -12,7 +12,7 @@ async fn main() -> Result<(), Error> {
let mut system = System::new().await?;
system.run().await?;
system.terminate().await?;
drop(system);
std::process::exit(0);
}
}

22
mantis/src/model/config.rs
View File

@ -21,9 +21,15 @@ pub struct SuricataConfig {
pub suppress: Vec<String>,
}
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()
}
}

1
mantis/src/model/error/ebpf.rs
View File

@ -70,4 +70,3 @@ traceable! {
FlowMapKeyMissing { direction: String, flow_direction: String, time_type: String } => tracing::Level::ERROR,
}
}

2
mantis/src/model/error/misc.rs
View File

@ -26,4 +26,4 @@ traceable! {
#[error("Failed to create traffic log file '{path}': {reason}")]
TrafficLogCreateError { path: String, reason: String } => tracing::Level::ERROR,
}
}
}

2
mantis/src/model/error/ml.rs
View File

@ -48,4 +48,4 @@ traceable! {
#[error("Failed to write CSV row: {reason}")]
TrafficLogWriteFailed { reason: String } => tracing::Level::WARN,
}
}
}

4
mantis/src/model/error/mod.rs
View File

@ -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<SuricataError> for Error {
fn from(error: SuricataError) -> Self {
Self::Suricata(error)
}
}
}

2
mantis/src/model/geo_stats.rs
View File

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

24
mantis/src/model/ip_address.rs
View File

@ -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())
}
}
}

4
mantis/src/model/log/ebpf.rs
View File

@ -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,
}
}
}

2
mantis/src/model/log/ml.rs
View File

@ -79,4 +79,4 @@ loggable! {
WindowDebug { src: String, pad: usize, window_size: usize, ae_score: f32, rows: String } => tracing::Level::DEBUG,
}
}
}

8
mantis/src/model/log/mod.rs
View File

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

2
mantis/src/model/log/suricata.rs
View File

@ -33,4 +33,4 @@ loggable! {
#[error("[suricata] {line}")]
ProcessError { line: String } => tracing::Level::ERROR,
}
}
}

2
mantis/src/model/log/system.rs
View File

@ -28,4 +28,4 @@ loggable! {
TrafficLoggingEnabled { path: String } => tracing::Level::INFO,
}
}
}

2
mantis/src/model/ml_detection.rs
View File

@ -236,4 +236,4 @@ impl UnifiedAlert {
rule_msg: Some(m.msg.clone()),
}
}
}
}

4
mantis/src/model/mod.rs
View File

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

9
mantis/src/utils/ip_address.rs
View File

@ -13,12 +13,7 @@ pub fn convert_ports_to_vec(ports: [u16; MAX_RULES_PORT]) -> Vec<Port> {
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()
}
}
}
}

6
mantis/src/utils/logging.rs
View File

@ -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");

8
mantis/src/utils/mod.rs
View File

@ -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;
pub mod ip_address;

2
mantis/src/utils/packet_parser.rs
View File

@ -184,4 +184,4 @@ pub fn format_ipv6(addr: u128) -> String {
bytes[14],
bytes[15]
)
}
}

26
mantis/src/web/api/auth.rs
View File

@ -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<AppState>,
Json(body): Json<LoginRequest>,
) -> impl IntoResponse {
async fn login(State(state): State<AppState>, Json(body): Json<LoginRequest>) -> 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(),

14
mantis/src/web/api/control/access_control.rs
View File

@ -63,7 +63,11 @@ async fn remove_ipv4_list(
State(state): State<AppState>,
Json(address): Json<SocketAddrV4>,
) -> 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<AppState>,
Json(address): Json<SocketAddrV6>,
) -> 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(),
}
}
}

76
mantis/src/web/api/control/service.rs
View File

@ -11,17 +11,57 @@ use crate::core::app_state::AppState;
pub fn router() -> Router<AppState> {
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<AppState>) -> impl IntoResponse {
@ -32,7 +72,10 @@ async fn get_ipv6_http_service(State(state): State<AppState>) -> impl IntoRespon
Json(state.service.get_ipv6_http_service().await)
}
async fn add_ipv4_http_service(State(state): State<AppState>, Json(payload): Json<(SocketAddrV4, Vec<HttpMethod>)>) -> impl IntoResponse {
async fn add_ipv4_http_service(
State(state): State<AppState>,
Json(payload): Json<(SocketAddrV4, Vec<HttpMethod>)>,
) -> 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<AppState>, Json(payload): Jso
}
}
async fn add_ipv6_http_service(State(state): State<AppState>, Json(payload): Json<(SocketAddrV6, Vec<HttpMethod>)>) -> impl IntoResponse {
async fn add_ipv6_http_service(
State(state): State<AppState>,
Json(payload): Json<(SocketAddrV6, Vec<HttpMethod>)>,
) -> 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<AppState>, Json(payload): Jso
}
}
async fn remove_ipv4_http_service(State(state): State<AppState>, Json(payload): Json<(SocketAddrV4, Vec<HttpMethod>)>) -> impl IntoResponse {
async fn remove_ipv4_http_service(
State(state): State<AppState>,
Json(payload): Json<(SocketAddrV4, Vec<HttpMethod>)>,
) -> 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<AppState>, Json(payload):
}
}
async fn remove_ipv6_http_service(State(state): State<AppState>, Json(payload): Json<(SocketAddrV6, Vec<HttpMethod>)>) -> impl IntoResponse {
async fn remove_ipv6_http_service(
State(state): State<AppState>,
Json(payload): Json<(SocketAddrV6, Vec<HttpMethod>)>,
) -> impl IntoResponse {
let (addr, methods) = payload;
match state.service.remove_ipv6_http_service(addr, methods).await {
Ok(_) => StatusCode::OK.into_response(),

34
mantis/src/web/api/control/statistics.rs
View File

@ -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<AppState> {
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<AppState>,
) -> 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<AppState>,
) -> 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<AppState>,
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<AppState>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
ws.on_upgrade(move |socket| {
flow_websocket::handle_ipv6_flow(socket, state, direction, flow_direction, time_type)
})
}
ws.on_upgrade(move |socket| flow_websocket::handle_ipv6_flow(socket, state, direction, flow_direction, time_type))
}

15
mantis/src/web/api/default.rs
View File

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

7
mantis/src/web/api/detection_alert.rs
View File

@ -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<AppState> {
Router::new().route("/websocket/alert", get(websocket_alert))
}
async fn websocket_alert(
ws: axum::extract::ws::WebSocketUpgrade,
State(state): State<AppState>,
) -> impl IntoResponse {
async fn websocket_alert(ws: axum::extract::ws::WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
let rx = state.detection_alert.subscribe();
ws.on_upgrade(|socket| alert_websocket::handle_alert(socket, rx))
}

4
mantis/src/web/api/health.rs
View File

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

4
mantis/src/web/api/misc.rs
View File

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

9
mantis/src/web/middleware/auth.rs
View File

@ -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<AppState> for AuthenticatedUser {
type Rejection = (StatusCode, &'static str);
async fn from_request_parts(
parts: &mut Parts,
state: &AppState,
) -> Result<Self, Self::Rejection> {
async fn from_request_parts(parts: &mut Parts, state: &AppState) -> Result<Self, Self::Rejection> {
let auth_cfg = state
.app_config
.auth

2
mantis/src/web/websocket/alert_websocket.rs
View File

@ -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<UnifiedAlert>) {
let (mut sender, mut receiver) = socket.split();

2
mantis/src/web/websocket/flow_websocket.rs
View File

@ -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};

2
mantis/src/web/websocket/health_websocket.rs
View File

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

2
mantis/src/web/websocket/mod.rs
View File

@ -1,3 +1,3 @@
pub mod alert_websocket;
pub mod flow_websocket;
pub mod health_websocket;
pub mod alert_websocket;