diff --git a/config.toml b/config.toml index 882413d..7a1a68c 100644 --- a/config.toml +++ b/config.toml @@ -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 \ No newline at end of file +ml_cpu = 7 + +ae_threshold_method = "94" \ No newline at end of file diff --git a/net-guardia/src/core/system.rs b/net-guardia/src/core/system.rs index 73f6a64..02c7b65 100644 --- a/net-guardia/src/core/system.rs +++ b/net-guardia/src/core/system.rs @@ -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(), diff --git a/net-guardia/src/detection/ml/config_loader.rs b/net-guardia/src/detection/ml/config_loader.rs index f70f699..b2e6646 100644 --- a/net-guardia/src/detection/ml/config_loader.rs +++ b/net-guardia/src/detection/ml/config_loader.rs @@ -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, pub ae_post_clip_min: f64, pub ae_post_clip_max: f64, - pub ae_threshold: f32, - pub ae_threshold_method: Option, + pub ae_thresholds: HashMap, 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 { + 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 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) } diff --git a/net-guardia/src/detection/ml/engine.rs b/net-guardia/src/detection/ml/engine.rs index b1d2ed8..f358496 100644 --- a/net-guardia/src/detection/ml/engine.rs +++ b/net-guardia/src/detection/ml/engine.rs @@ -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 = 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 = 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 = 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); } } diff --git a/net-guardia/src/detection/ml/feature_extractor.rs b/net-guardia/src/detection/ml/feature_extractor.rs index 98bdf66..1344681 100644 --- a/net-guardia/src/detection/ml/feature_extractor.rs +++ b/net-guardia/src/detection/ml/feature_extractor.rs @@ -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() -} +} \ No newline at end of file diff --git a/net-guardia/src/detection/ml/flow_tracker.rs b/net-guardia/src/detection/ml/flow_tracker.rs index 76684a3..15ff66f 100644 --- a/net-guardia/src/detection/ml/flow_tracker.rs +++ b/net-guardia/src/detection/ml/flow_tracker.rs @@ -311,14 +311,6 @@ impl FlowTracker { self.flows.values().cloned().collect() } - pub fn get_flows_for_inference(&self, min_packets: usize) -> Vec { - 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 -} +} \ No newline at end of file diff --git a/net-guardia/src/detection/ml/inference.rs b/net-guardia/src/detection/ml/inference.rs index 02f2c7e..227e731 100644 --- a/net-guardia/src/detection/ml/inference.rs +++ b/net-guardia/src/detection/ml/inference.rs @@ -18,15 +18,24 @@ use crate::model::ml_detection::DetectionResult; pub struct Inference { pub models: Arc, pub config: Arc, + 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>>>, } impl Inference { pub fn new(models: Arc, config: Arc) -> 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::>() + + // 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> = (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::>().join(" "); + format!(" t{t:02}: [{vals}]") + } + }).collect::>().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) } -} +} \ No newline at end of file diff --git a/net-guardia/src/model/config.rs b/net-guardia/src/model/config.rs index 2479917..4b5a4ea 100644 --- a/net-guardia/src/model/config.rs +++ b/net-guardia/src/model/config.rs @@ -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() } \ No newline at end of file diff --git a/net-guardia/src/model/error/ml.rs b/net-guardia/src/model/error/ml.rs index 21806a4..a22bb4b 100644 --- a/net-guardia/src/model/error/ml.rs +++ b/net-guardia/src/model/error/ml.rs @@ -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, diff --git a/net-guardia/src/model/log/ml.rs b/net-guardia/src/model/log/ml.rs index 5ecad54..0709fb3 100644 --- a/net-guardia/src/model/log/ml.rs +++ b/net-guardia/src/model/log/ml.rs @@ -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, + } } \ No newline at end of file diff --git a/net-guardia/static/artifacts/deep_autoencoder.onnx b/net-guardia/static/artifacts/deep_autoencoder.onnx index 8275cb1..831e6d9 100644 Binary files a/net-guardia/static/artifacts/deep_autoencoder.onnx and b/net-guardia/static/artifacts/deep_autoencoder.onnx differ diff --git a/net-guardia/static/artifacts/full_config.json b/net-guardia/static/artifacts/full_config.json index fecedda..2aa5a31 100644 --- a/net-guardia/static/artifacts/full_config.json +++ b/net-guardia/static/artifacts/full_config.json @@ -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": {