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: require full window fill before ONNX inference

Browse Source
This commit is contained in:
ParrotXray 2026-05-20 06:23:08 +00:00
parent 0f2b146bb3
commit 3a17e62dcd
12 changed files with 146 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
config.toml
View File

@ -41,4 +41,6 @@ xsk_cpu_base = 0
#
# Pin ML inference (ONNX spawn_blocking) to this core.
# Example: on an 8-core machine, reserve core 7 for inference.
ml_cpu = 7
ml_cpu = 7
ae_threshold_method = "94"

2
net-guardia/src/core/system.rs
View File

@ -41,7 +41,7 @@ impl System {
let (mut egress_ebpf, egress_program_array) = System::get_egress_ebpf()?;
let app_config = Arc::new(AppConfig::new()?);
let inference_config = Arc::new(InferenceConfig::load_file(&app_config.models_config_name)?);
let inference_config = Arc::new(InferenceConfig::load_file(&app_config.models_config_name, &app_config.ae_threshold_method)?);
let ebpf_services = Arc::new(EbpfServices::new(
app_config.clone(),

19
net-guardia/src/detection/ml/config_loader.rs
View File

@ -3,6 +3,7 @@ use std::fs;
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
use macros::log;
use crate::model::error::ml::MLError;
use crate::model::ml_detection::ClipParams;
@ -14,18 +15,28 @@ pub struct InferenceConfig {
pub ae_scaler_std: Vec<f64>,
pub ae_post_clip_min: f64,
pub ae_post_clip_max: f64,
pub ae_threshold: f32,
pub ae_threshold_method: Option<String>,
pub ae_thresholds: HashMap<String, f64>,
pub window_size: usize,
/// Selected threshold method name, set at load time from config.toml.
#[serde(skip)]
pub ae_threshold_method: String,
}
impl InferenceConfig {
pub fn load_file(file: &str) -> Result<Self, MLError> {
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 config: InferenceConfig = serde_json::from_str(&content)
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() });
}
config.ae_threshold_method = method.to_string();
Ok(config)
}

46
net-guardia/src/detection/ml/engine.rs
View File

@ -90,7 +90,7 @@ impl Engine {
_ = ticker.tick() => {}
}
let (total_flows, packet_counts, flows) = {
let (total_flows, packet_counts, flows, active_ips) = {
let Ok(mut t) = self.tracker.lock() else {
log!(MLError::TrackerLockPoisoned);
continue;
@ -100,12 +100,13 @@ impl Engine {
.iter()
.map(|f| f.packet_count())
.collect();
let flows = if self.traffic_logger.is_some() {
t.drain_flows(self.min_packets)
} else {
t.get_flows_for_inference(self.min_packets)
};
(total_flows, packet_counts, flows)
let flows = t.drain_flows(self.min_packets);
t.cleanup_old_flows(self.flow_timeout_us);
let active_ips: std::collections::HashSet<String> = t
.get_flows_snapshot().iter()
.map(|f| f.flow_key.src_ip.clone())
.collect();
(total_flows, packet_counts, flows, active_ips)
};
log!(MLLog::FlowStats(
@ -116,19 +117,7 @@ impl Engine {
));
if flows.is_empty() {
if self.traffic_logger.is_none() {
if let Ok(mut t) = self.tracker.lock() {
t.cleanup_old_flows(self.flow_timeout_us);
let active_ips: std::collections::HashSet<String> = t
.get_flows_snapshot().iter()
.map(|f| f.flow_key.src_ip.clone())
.collect();
drop(t);
self.inference_pipeline.cleanup_buffers(&active_ips);
} else {
log!(MLError::TrackerLockPoisoned);
}
}
self.inference_pipeline.cleanup_buffers(&active_ips);
continue;
}
@ -138,11 +127,6 @@ impl Engine {
let features = FlowFeatures::extract(flow, &feature_names);
logger.log_row(features.to_csv_record());
}
if let Ok(mut t) = self.tracker.lock() {
t.cleanup_old_flows(self.flow_timeout_us);
} else {
log!(MLError::TrackerLockPoisoned);
}
continue;
}
@ -209,17 +193,7 @@ impl Engine {
aggregator.cleanup();
}
if let Ok(mut t) = self.tracker.lock() {
t.cleanup_old_flows(self.flow_timeout_us);
let active_ips: std::collections::HashSet<String> = t
.get_flows_snapshot().iter()
.map(|f| f.flow_key.src_ip.clone())
.collect();
drop(t);
self.inference_pipeline.cleanup_buffers(&active_ips);
} else {
log!(MLError::TrackerLockPoisoned);
}
self.inference_pipeline.cleanup_buffers(&active_ips);
}
}

106
net-guardia/src/detection/ml/feature_extractor.rs
View File

@ -95,72 +95,72 @@ impl FlowFeatures {
"Total Backward Packets" | "Tot Bwd Pkts" | "bwd_packets" => bwd_count,
"Total Length of Fwd Packets" | "TotLen Fwd Pkts" | "fwd_bytes" => flow.fwd_total_bytes as f64,
"Total Length of Bwd Packets" | "TotLen Bwd Pkts" | "bwd_bytes" => flow.bwd_total_bytes as f64,
"Fwd Packet Length Max" => fwd_max,
"Fwd Packet Length Min" => fwd_min,
"Fwd Packet Length Max" | "fwd_pkt_len_max" => fwd_max,
"Fwd Packet Length Min" | "fwd_pkt_len_min" => fwd_min,
"Fwd Packet Length Mean" | "Fwd Pkt Len Mean" | "fwd_pkt_len_mean" => fwd_mean,
"Fwd Packet Length Std" | "Fwd Pkt Len Std" | "fwd_pkt_len_std" => fwd_std,
"Bwd Packet Length Max" => bwd_max,
"Bwd Packet Length Min" => bwd_min,
"Bwd Packet Length Max" | "bwd_pkt_len_max" => bwd_max,
"Bwd Packet Length Min" | "bwd_pkt_len_min" => bwd_min,
"Bwd Packet Length Mean" | "Bwd Pkt Len Mean" | "bwd_pkt_len_mean" => bwd_mean,
"Bwd Packet Length Std" | "Bwd Pkt Len Std" | "bwd_pkt_len_std" => bwd_std,
"Flow Bytes/s" | "Flow Byts/s" | "flow_bytes_per_sec" => safe_div(total_bytes, duration_s),
"Flow Packets/s" | "Flow Pkts/s" | "flow_pkts_per_sec" => safe_div(total_count, duration_s),
"Flow IAT Mean" | "flow_iat_mean" => flow_iat_mean,
"Flow IAT Std" => flow_iat_std,
"Flow IAT Max" => flow_iat_max,
"Flow IAT Min" => flow_iat_min,
"Fwd IAT Total" => fwd_iat_total,
"Flow IAT Std" | "flow_iat_std" => flow_iat_std,
"Flow IAT Max" | "flow_iat_max" => flow_iat_max,
"Flow IAT Min" | "flow_iat_min" => flow_iat_min,
"Fwd IAT Total" | "fwd_iat_total" => fwd_iat_total,
"Fwd IAT Mean" | "fwd_iat_mean" => fwd_iat_mean,
"Fwd IAT Std" => fwd_iat_std,
"Fwd IAT Max" => fwd_iat_max,
"Fwd IAT Min" => fwd_iat_min,
"Bwd IAT Total" => bwd_iat_total,
"Fwd IAT Std" | "fwd_iat_std" => fwd_iat_std,
"Fwd IAT Max" | "fwd_iat_max" => fwd_iat_max,
"Fwd IAT Min" | "fwd_iat_min" => fwd_iat_min,
"Bwd IAT Total" | "bwd_iat_total" => bwd_iat_total,
"Bwd IAT Mean" | "bwd_iat_mean" => bwd_iat_mean,
"Bwd IAT Std" => bwd_iat_std,
"Bwd IAT Max" => bwd_iat_max,
"Bwd IAT Min" => bwd_iat_min,
"Fwd PSH Flags" => fwd_psh,
"Bwd PSH Flags" => bwd_psh,
"Fwd URG Flags" => fwd_urg,
"Bwd URG Flags" => bwd_urg,
"Fwd Header Length" => flow.fwd_header_bytes as f64,
"Bwd Header Length" => flow.bwd_header_bytes as f64,
"Fwd Packets/s" => safe_div(fwd_count, duration_s),
"Bwd Packets/s" => safe_div(bwd_count, duration_s),
"Min Packet Length" => min_len,
"Max Packet Length" => max_len,
"Bwd IAT Std" | "bwd_iat_std" => bwd_iat_std,
"Bwd IAT Max" | "bwd_iat_max" => bwd_iat_max,
"Bwd IAT Min" | "bwd_iat_min" => bwd_iat_min,
"Fwd PSH Flags" | "fwd_psh_flags" => fwd_psh,
"Bwd PSH Flags" | "bwd_psh_flags" => bwd_psh,
"Fwd URG Flags" | "fwd_urg_flags" => fwd_urg,
"Bwd URG Flags" | "bwd_urg_flags" => bwd_urg,
"Fwd Header Length" | "fwd_header_length" => flow.fwd_header_bytes as f64,
"Bwd Header Length" | "bwd_header_length" => flow.bwd_header_bytes as f64,
"Fwd Packets/s" | "fwd_pkts_per_sec" => safe_div(fwd_count, duration_s),
"Bwd Packets/s" | "bwd_pkts_per_sec" => safe_div(bwd_count, duration_s),
"Min Packet Length" | "pkt_len_min" => min_len,
"Max Packet Length" | "pkt_len_max" => max_len,
"Packet Length Mean" | "Pkt Len Mean" | "pkt_len_mean" => mean_len,
"Packet Length Std" | "Pkt Len Std" | "pkt_len_std" => std_len,
"Packet Length Variance" => std_len * std_len,
"Packet Length Variance" | "pkt_len_var" => std_len * std_len,
"FIN Flag Count" | "FIN Flag Cnt" | "fin_flag_cnt" => flow.fin_count as f64,
"SYN Flag Count" | "SYN Flag Cnt" | "syn_flag_cnt" => flow.syn_count as f64,
"RST Flag Count" | "RST Flag Cnt" | "rst_flag_cnt" => flow.rst_count as f64,
"PSH Flag Count" | "PSH Flag Cnt" | "psh_flag_cnt" => flow.psh_count as f64,
"ACK Flag Count" | "ACK Flag Cnt" | "ack_flag_cnt" => flow.ack_count as f64,
"URG Flag Count" => flow.urg_count as f64,
"CWE Flag Count" => flow.cwe_count as f64,
"ECE Flag Count" => flow.ece_count as f64,
"Down/Up Ratio" => safe_div(bwd_count, fwd_count),
"Average Packet Size" => safe_div(total_bytes, total_count),
"Avg Fwd Segment Size" => safe_div(flow.fwd_total_bytes as f64, fwd_count),
"Avg Bwd Segment Size" => safe_div(flow.bwd_total_bytes as f64, bwd_count),
"Fwd Header Length.1" => flow.fwd_header_bytes as f64,
"Fwd Avg Bytes/Bulk" => safe_div(fwd_bulk.total_bytes as f64, fwd_bulk.bulk_count as f64),
"Fwd Avg Packets/Bulk" => safe_div(fwd_bulk.total_packets as f64, fwd_bulk.bulk_count as f64),
"Fwd Avg Bulk Rate" => safe_div(
"URG Flag Count" | "urg_flag_cnt" => flow.urg_count as f64,
"CWE Flag Count" | "cwe_flag_cnt" => flow.cwe_count as f64,
"ECE Flag Count" | "ece_flag_cnt" => flow.ece_count as f64,
"Down/Up Ratio" | "down_up_ratio" => safe_div(bwd_count, fwd_count),
"Average Packet Size" | "avg_pkt_size" => safe_div(total_bytes, total_count),
"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 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" => safe_div(bwd_bulk.total_bytes as f64, bwd_bulk.bulk_count as f64),
"Bwd Avg Packets/Bulk" => safe_div(bwd_bulk.total_packets as f64, bwd_bulk.bulk_count as f64),
"Bwd Avg Bulk Rate" => safe_div(
"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,
),
"Subflow Fwd Packets" => fwd_count,
"Subflow Fwd Bytes" => flow.fwd_total_bytes as f64,
"Subflow Bwd Packets" => bwd_count,
"Subflow Bwd Bytes" => flow.bwd_total_bytes as f64,
"Subflow Fwd Packets" | "subflow_fwd_packets" => fwd_count,
"Subflow Fwd Bytes" | "subflow_fwd_bytes" => flow.fwd_total_bytes as f64,
"Subflow Bwd Packets" | "subflow_bwd_packets" => bwd_count,
"Subflow Bwd Bytes" | "subflow_bwd_bytes" => flow.bwd_total_bytes as f64,
"Init_Win_bytes_forward" | "Init Fwd Win Byts" | "fwd_win_bytes" => flow.init_win_bytes_fwd as f64,
"Init_Win_bytes_backward" | "Init Bwd Win Byts" | "bwd_win_bytes" => flow.init_win_bytes_bwd as f64,
"act_data_pkt_fwd" | "Fwd Act Data Pkts" | "fwd_act_data_pkts" => flow.act_data_pkt_fwd as f64,
@ -169,14 +169,14 @@ impl FlowFeatures {
.min_by(|a, b| a.total_cmp(b))
.copied()
.unwrap_or(0.0),
"Active Mean" => active_mean,
"Active Std" => active_std,
"Active Max" => active_max,
"Active Min" => active_min,
"Idle Mean" => idle_mean,
"Idle Std" => idle_std,
"Idle Max" => idle_max,
"Idle Min" => idle_min,
"Active Mean" | "active_mean" => active_mean,
"Active Std" | "active_std" => active_std,
"Active Max" | "active_max" => active_max,
"Active Min" | "active_min" => active_min,
"Idle Mean" | "idle_mean" => idle_mean,
"Idle Std" | "idle_std" => idle_std,
"Idle Max" | "idle_max" => idle_max,
"Idle Min" | "idle_min" => idle_min,
_ => 0.0,
}
}
@ -358,4 +358,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()
}
}

10
net-guardia/src/detection/ml/flow_tracker.rs
View File

@ -311,14 +311,6 @@ impl FlowTracker {
self.flows.values().cloned().collect()
}
pub fn get_flows_for_inference(&self, min_packets: usize) -> Vec<FlowData> {
self.flows
.values()
.filter(|flow| flow.packet_count() >= min_packets)
.cloned()
.collect()
}
pub fn cleanup_old_flows(&mut self, max_age_us: u64) {
let now = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
@ -378,4 +370,4 @@ fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16)
}
None
}
}

55
net-guardia/src/detection/ml/inference.rs
View File

@ -18,15 +18,24 @@ use crate::model::ml_detection::DetectionResult;
pub struct Inference {
pub models: Arc<MLModels>,
pub config: Arc<InferenceConfig>,
pub threshold: f32,
// Must equal window_size: only a fully-filled buffer is in-distribution.
// Zero-padded sequences are OOD for the autoencoder; any pad > 0 produces
// inflated MSE regardless of traffic type, causing false positives.
min_window_fill: usize,
// per-src_ip sliding window buffer: src_ip -> deque of feature vectors
flow_buffers: Mutex<HashMap<String, VecDeque<Vec<f32>>>>,
}
impl Inference {
pub fn new(models: Arc<MLModels>, config: Arc<InferenceConfig>) -> Self {
let threshold = config.ae_thresholds[&config.ae_threshold_method] as f32;
let min_window_fill = config.window_size;
Self {
models,
config,
threshold,
min_window_fill,
flow_buffers: Mutex::new(HashMap::new()),
}
}
@ -49,7 +58,9 @@ impl Inference {
let t1 = Instant::now();
let window_size = self.config.window_size;
let sequence = {
let feat_len = self.config.num_ae_features();
let (sequence, buf_len_snapshot) = {
let Ok(mut buffers) = self.flow_buffers.lock() else {
log!(MLError::InferenceLockPoisoned);
return None;
@ -62,14 +73,23 @@ impl Inference {
if buf.len() > window_size {
buf.pop_front();
}
if buf.len() < window_size {
return None;
}
buf.iter().cloned().collect::<Vec<_>>()
// Zero-pad on the left (scaled space) when fewer than window_size
// flows have been seen for this src_ip — matches Python's
// _make_per_flow_sequences which initialises sequences to zeros
// before filling from the right.
let buf_len = buf.len();
let pad = window_size - buf_len;
let mut seq: Vec<Vec<f32>> = (0..pad).map(|_| vec![0.0f32; feat_len]).collect();
seq.extend(buf.iter().cloned());
(seq, buf_len)
};
if buf_len_snapshot < self.min_window_fill {
return None;
}
let t2 = Instant::now();
let feat_len = self.config.num_ae_features();
let ae_input = Array3::from_shape_fn(
(1, window_size, feat_len),
|(_, t, f)| sequence[t][f],
@ -93,7 +113,24 @@ impl Inference {
t4.duration_since(t3).as_millis() as u64,
));
let is_attack = ae_score >= self.config.ae_threshold;
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");
log!(MLLog::WindowDebug(
flow.flow_key.src_ip.clone(),
pad,
window_size,
ae_score,
rows,
));
let is_attack = ae_score >= self.threshold;
let flow_key = format!(
"{}:{} -> {}:{} (proto {}) [{}]",
flow.flow_key.src_ip,
@ -112,7 +149,7 @@ impl Inference {
attack_type: if is_attack { Some("ANOMALY".to_string()) } else { None },
confidence: ae_score,
ae_score,
threshold: self.config.ae_threshold,
threshold: self.threshold,
})
}
@ -143,4 +180,4 @@ impl Inference {
Ok(mse)
}
}
}

9
net-guardia/src/model/config.rs
View File

@ -52,6 +52,11 @@ pub struct Config {
/// Only used in "or" (corroboration window) and "and" modes. Defaults to 10.
#[serde(default = "default_fusion_window_secs")]
pub fusion_window_secs: u64,
/// Key into ae_thresholds in inference_config.json that selects the active
/// anomaly detection threshold. Valid values: "90".."99", "mean+2std",
/// "mean+1std", "Q3+1.5IQR", "Q3+3.0IQR". Defaults to "95" when absent.
#[serde(default = "default_ae_threshold_method")]
pub ae_threshold_method: String,
}
fn default_fusion_mode() -> String {
@ -60,4 +65,8 @@ fn default_fusion_mode() -> String {
fn default_fusion_window_secs() -> u64 {
10
}
fn default_ae_threshold_method() -> String {
"95".to_string()
}

4
net-guardia/src/model/error/ml.rs
View File

@ -20,6 +20,10 @@ traceable! {
#[error("Failed to parse inference configuration: {reason}")]
ConfigParseFailed { reason: String } => tracing::Level::ERROR,
#[no_source]
#[error("Threshold method '{method}' not found in ae_thresholds map")]
ThresholdMethodNotFound { method: String } => tracing::Level::ERROR,
#[no_source]
#[error("Tracker mutex poisoned; skipping operation")]
TrackerLockPoisoned => tracing::Level::ERROR,

3
net-guardia/src/model/log/ml.rs
View File

@ -81,5 +81,8 @@ loggable! {
#[error("Timing [{src}]: feature={feature_ms}ms buffer={buffer_ms}ms tensor={tensor_ms}ms onnx={onnx_ms}ms")]
InferenceTiming { src: String, feature_ms: u64, buffer_ms: u64, tensor_ms: u64, onnx_ms: u64 } => tracing::Level::INFO,
#[error("Window [{src}] pad={pad}/{window_size} ae={ae_score:.6}\n{rows}")]
WindowDebug { src: String, pad: usize, window_size: usize, ae_score: f32, rows: String } => tracing::Level::DEBUG,
}
}

BIN
net-guardia/static/artifacts/deep_autoencoder.onnx
View File

Binary file not shown.

2
net-guardia/static/artifacts/full_config.json
View File

@ -1,5 +1,5 @@
{
"created_at": "2026-05-07T07:34:08.124387",
"created_at": "2026-05-20T09:38:52.941433",
"framework": "PyTorch",
"model": {
"lstm_deep_autoencoder": {