mirror of
https://github.com/DaLaw2/NetGuardia.git
synced 2026-08-24 14:10:28 +09:00
feat(v12-m1): BYO-model foundation — manifest loader + FEATURE_REGISTRY
Land v12 M1. Promote the 88-arm feature match in PrecomputedStats::get()
to a FEATURE_REGISTRY (name → fn pointer) keyed by every legacy alias;
manifest validation consults it via feature_is_known(). Add a YAML
ModelManifest parser with AdapterKind { classifier_only, autoencoder_only,
multi_task }, per-label metadata (confirmations/playbook — parsed, wired
in M2), thresholds, and a scaler sidecar pointer. Introspect ONNX input
shape via tract's InferenceFact and cross-check against the manifest-
declared feature count; mismatches surface as MLError::FeatureMismatch
with both declared and onnx dims. Compose runtime MLInferenceConfig from
(manifest + scaler sidecar JSON) with order-sensitive feature reconcili-
ation — any drift between the two is rejected, not silently overridden.
service_factory prefers models/manifest.yaml when present; the legacy
JSON-only path remains as fallback. For multi_task manifests, model
filenames come from the manifest, not AppConfig. Ship models/manifest.yaml
describing the current v10 multi_task model (31 AE features, 10 labels).
Incidental clippy fixes bundled: reorder mod tests / impl blocks in
adapter/http/setup.rs and infrastructure/secret_store.rs to silence
items_after_test_module under --tests; rustfmt polish on unrelated drift.
New error variants: ManifestInvalid, FeatureMismatch, UnknownFeature.
New log variants: ManifestLoaded, OnnxShapeChecked.
Tests: +13 ML-layer tests — registry coverage of v10 features, alias
consistency, safe_div zero-denominator, manifest parse/validate, sidecar
feature-count rejection, real v10 ONNX introspection (AE=31,
classifier=32), end-to-end load_named via manifest paths. Suite 157/157;
cargo clippy --package net-guardia --tests -- -D warnings clean.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
c6a77c9611
commit
86358074b9
20
Cargo.lock
generated
20
Cargo.lock
generated
@ -2381,6 +2381,7 @@ dependencies = [
|
||||
"sd-notify",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml_ng",
|
||||
"sha2",
|
||||
"sysinfo",
|
||||
"thiserror 2.0.18",
|
||||
@ -3437,6 +3438,19 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_yaml_ng"
|
||||
version = "0.10.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
"unsafe-libyaml",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha1"
|
||||
version = "0.10.6"
|
||||
@ -4213,6 +4227,12 @@ dependencies = [
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unsafe-libyaml"
|
||||
version = "0.2.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.9.0"
|
||||
|
||||
@ -20,6 +20,7 @@ libxdp-sys = { version = "0.2.4", features = ["use_cc_build", "use_precompiled_b
|
||||
# Serialization
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = "1.0.149"
|
||||
serde_yaml_ng = "0.10.0"
|
||||
|
||||
# Async runtime
|
||||
tokio = { version = "1.50.0", features = ["rt-multi-thread", "macros", "sync", "time", "process", "io-util", "fs", "signal"] }
|
||||
|
||||
69
models/manifest.yaml
Normal file
69
models/manifest.yaml
Normal file
@ -0,0 +1,69 @@
|
||||
# NetGuardia v10 model manifest — authored 2026-04-16 as part of v12 M1.
|
||||
# Schema: structural/semantic info lives here; preprocessing arrays
|
||||
# (scaler mean/std, clip params, feature weights) stay in inference_config.json
|
||||
# referenced by `preprocessing.scaler_sidecar` below.
|
||||
|
||||
name: netguardia-v10
|
||||
adapter: multi_task
|
||||
|
||||
models:
|
||||
autoencoder: deep_autoencoder.onnx
|
||||
classifier: classifier.onnx
|
||||
|
||||
# 31 AE-input features. Order matters — must match ONNX input column order
|
||||
# and inference_config.json `ae_feature_names`. The classifier takes these
|
||||
# plus `ae_anomaly_score` appended as the 32nd input (handled in code).
|
||||
features:
|
||||
- flow_duration
|
||||
- fwd_packets
|
||||
- bwd_packets
|
||||
- fwd_bytes
|
||||
- bwd_bytes
|
||||
- flow_bytes_per_sec
|
||||
- flow_pkts_per_sec
|
||||
- fwd_win_bytes
|
||||
- bwd_win_bytes
|
||||
- fwd_pkt_len_mean
|
||||
- bwd_pkt_len_mean
|
||||
- fwd_iat_mean
|
||||
- bwd_iat_mean
|
||||
- flow_iat_mean
|
||||
- pkt_len_mean
|
||||
- dst_port
|
||||
- protocol
|
||||
- psh_flag_cnt
|
||||
- ack_flag_cnt
|
||||
- syn_flag_cnt
|
||||
- fin_flag_cnt
|
||||
- rst_flag_cnt
|
||||
- pkt_len_std
|
||||
- fwd_pkt_len_std
|
||||
- bwd_pkt_len_std
|
||||
- fwd_seg_size_min
|
||||
- fwd_act_data_pkts
|
||||
- fwd_iat_std
|
||||
- bwd_iat_std
|
||||
- fwd_bwd_bytes_ratio
|
||||
- iat_cv
|
||||
|
||||
# `confirmations` / `playbook` are parsed today; M2 wires them into the aggregator.
|
||||
labels:
|
||||
"0": { name: Bot }
|
||||
"1": { name: Brute Force }
|
||||
"2": { name: C2 Communication, confirmations: 1 }
|
||||
"3": { name: DNS Tunneling, confirmations: 1 }
|
||||
"4": { name: DoS/DDoS, confirmations: 2 }
|
||||
"5": { name: Exploitation, confirmations: 1 }
|
||||
"6": { name: Malware }
|
||||
"7": { name: Normal }
|
||||
"8": { name: Reconnaissance }
|
||||
"9": { name: Web Attack }
|
||||
|
||||
thresholds:
|
||||
anomaly: 0.9179317355155945
|
||||
c2: 0.9085615873336792
|
||||
class_min_confidence: 0.4
|
||||
ae: 0.23011694848537445
|
||||
|
||||
preprocessing:
|
||||
scaler_sidecar: inference_config.json
|
||||
@ -28,6 +28,7 @@ tokio-tungstenite = "0.28.0"
|
||||
# Serialization
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml_ng = { workspace = true }
|
||||
toml = "1.0.7"
|
||||
|
||||
# Async
|
||||
|
||||
@ -172,47 +172,6 @@ async fn complete_setup(
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_valid_interface_names() {
|
||||
assert!(is_valid_interface_name("eth0"));
|
||||
assert!(is_valid_interface_name("ens33"));
|
||||
assert!(is_valid_interface_name("br-lan"));
|
||||
assert!(is_valid_interface_name("wlan0.1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_empty() {
|
||||
assert!(!is_valid_interface_name(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_too_long() {
|
||||
let long = "a".repeat(17);
|
||||
assert!(!is_valid_interface_name(&long));
|
||||
// Exactly 16 should be valid
|
||||
let exact = "a".repeat(16);
|
||||
assert!(is_valid_interface_name(&exact));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_path_traversal() {
|
||||
assert!(!is_valid_interface_name("../etc"));
|
||||
assert!(!is_valid_interface_name("../../shadow"));
|
||||
assert!(!is_valid_interface_name("/sys/class"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_special_chars() {
|
||||
assert!(!is_valid_interface_name("eth0;rm"));
|
||||
assert!(!is_valid_interface_name("lo&&cat"));
|
||||
assert!(!is_valid_interface_name("eth0 space"));
|
||||
}
|
||||
}
|
||||
|
||||
fn save_config(
|
||||
db: &Database,
|
||||
secrets: &dyn SecretStorePort,
|
||||
@ -258,3 +217,44 @@ fn save_config(
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_valid_interface_names() {
|
||||
assert!(is_valid_interface_name("eth0"));
|
||||
assert!(is_valid_interface_name("ens33"));
|
||||
assert!(is_valid_interface_name("br-lan"));
|
||||
assert!(is_valid_interface_name("wlan0.1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_empty() {
|
||||
assert!(!is_valid_interface_name(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_too_long() {
|
||||
let long = "a".repeat(17);
|
||||
assert!(!is_valid_interface_name(&long));
|
||||
// Exactly 16 should be valid
|
||||
let exact = "a".repeat(16);
|
||||
assert!(is_valid_interface_name(&exact));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_path_traversal() {
|
||||
assert!(!is_valid_interface_name("../etc"));
|
||||
assert!(!is_valid_interface_name("../../shadow"));
|
||||
assert!(!is_valid_interface_name("/sys/class"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_interface_special_chars() {
|
||||
assert!(!is_valid_interface_name("eth0;rm"));
|
||||
assert!(!is_valid_interface_name("lo&&cat"));
|
||||
assert!(!is_valid_interface_name("eth0 space"));
|
||||
}
|
||||
}
|
||||
|
||||
@ -1615,11 +1615,9 @@ impl Database {
|
||||
|
||||
let tx = conn.transaction()?;
|
||||
let prev_hash: String = tx
|
||||
.query_row(
|
||||
"SELECT row_hash FROM audit_log ORDER BY id DESC LIMIT 1",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.query_row("SELECT row_hash FROM audit_log ORDER BY id DESC LIMIT 1", [], |row| {
|
||||
row.get(0)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let row_hash = Self::audit_row_hash(&ts, actor, action, detail, &prev_hash);
|
||||
|
||||
@ -306,10 +306,7 @@ impl<T: NativeConvert + Pod> EntryMap<T> {
|
||||
let Some(map) = self.map.as_ref() else {
|
||||
return Vec::new();
|
||||
};
|
||||
map.keys()
|
||||
.filter_map(Result::ok)
|
||||
.map(|key| key.into_native())
|
||||
.collect()
|
||||
map.keys().filter_map(Result::ok).map(|key| key.into_native()).collect()
|
||||
}
|
||||
|
||||
fn add(&mut self, key: T::Native) -> Result<(), Error> {
|
||||
|
||||
@ -1,29 +1,199 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use super::manifest::{LabelSpec, ModelManifest};
|
||||
use crate::model::error::ml::MLError;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
|
||||
impl MLInferenceConfig {
|
||||
pub fn load_file(file: &str) -> Result<Self, MLError> {
|
||||
let path = PathBuf::from("models").join(file);
|
||||
let content = fs::read_to_string(&path).map_err(|_| MLError::ConfigLoadFailed(path.to_path_buf()))?;
|
||||
Self::load_file_at(&path)
|
||||
}
|
||||
|
||||
fn load_file_at(path: &Path) -> Result<Self, MLError> {
|
||||
let content = fs::read_to_string(path).map_err(|_| MLError::ConfigLoadFailed(path.to_path_buf()))?;
|
||||
let config: MLInferenceConfig =
|
||||
serde_json::from_str(&content).map_err(|e| MLError::ConfigParseFailed(e.to_string()))?;
|
||||
if config.ae_feature_names.is_empty() {
|
||||
return Err(MLError::ConfigParseFailed("ae_feature_names is empty"));
|
||||
}
|
||||
if config.ae_scaler_mean.len() != config.ae_feature_names.len() {
|
||||
return Err(MLError::ConfigParseFailed("scaler mean length mismatch"));
|
||||
}
|
||||
if config.ae_scaler_std.len() != config.ae_feature_names.len() {
|
||||
return Err(MLError::ConfigParseFailed("scaler std length mismatch"));
|
||||
}
|
||||
if config.output_names.len() != 3 {
|
||||
return Err(MLError::ConfigParseFailed(
|
||||
"MultiTaskModel requires exactly 3 output_names (anomaly, class_probs, c2_score)",
|
||||
));
|
||||
}
|
||||
validate(&config)?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Load a `MLInferenceConfig` by combining a manifest with its scaler sidecar.
|
||||
/// The manifest supplies the authoritative feature list and per-label metadata;
|
||||
/// the sidecar JSON supplies the numeric preprocessing arrays (scaler / clip / weights).
|
||||
///
|
||||
/// Consistency is enforced: manifest `features` must match the sidecar's
|
||||
/// `ae_feature_names` exactly (order-sensitive). Any drift between the two is a
|
||||
/// deployment bug, not a silent override.
|
||||
pub fn from_manifest_with_sidecar(manifest_path: &Path) -> Result<(Self, ModelManifest), MLError> {
|
||||
let manifest = ModelManifest::load(manifest_path)?;
|
||||
let sidecar_rel = manifest
|
||||
.preprocessing
|
||||
.as_ref()
|
||||
.map(|p| p.scaler_sidecar.as_str())
|
||||
.ok_or_else(|| {
|
||||
MLError::ManifestInvalid(
|
||||
manifest_path.to_path_buf(),
|
||||
"preprocessing.scaler_sidecar is required for v1 (scaler arrays live there)".to_string(),
|
||||
)
|
||||
})?;
|
||||
let sidecar_path = ModelManifest::resolve_relative(manifest_path, sidecar_rel);
|
||||
let mut config = Self::load_file_at(&sidecar_path)?;
|
||||
|
||||
reconcile_features(&manifest, &config, manifest_path)?;
|
||||
apply_manifest_overrides(&manifest, &mut config);
|
||||
|
||||
Ok((config, manifest))
|
||||
}
|
||||
}
|
||||
|
||||
fn validate(config: &MLInferenceConfig) -> Result<(), MLError> {
|
||||
if config.ae_feature_names.is_empty() {
|
||||
return Err(MLError::ConfigParseFailed("ae_feature_names is empty"));
|
||||
}
|
||||
if config.ae_scaler_mean.len() != config.ae_feature_names.len() {
|
||||
return Err(MLError::ConfigParseFailed("scaler mean length mismatch"));
|
||||
}
|
||||
if config.ae_scaler_std.len() != config.ae_feature_names.len() {
|
||||
return Err(MLError::ConfigParseFailed("scaler std length mismatch"));
|
||||
}
|
||||
if config.output_names.len() != 3 {
|
||||
return Err(MLError::ConfigParseFailed(
|
||||
"MultiTaskModel requires exactly 3 output_names (anomaly, class_probs, c2_score)",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reconcile_features(
|
||||
manifest: &ModelManifest,
|
||||
config: &MLInferenceConfig,
|
||||
manifest_path: &Path,
|
||||
) -> Result<(), MLError> {
|
||||
if manifest.features.len() != config.ae_feature_names.len() {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.to_path_buf(),
|
||||
format!(
|
||||
"feature count mismatch with sidecar: manifest declares {}, sidecar lists {}",
|
||||
manifest.features.len(),
|
||||
config.ae_feature_names.len()
|
||||
),
|
||||
));
|
||||
}
|
||||
for (i, (mf, sf)) in manifest.features.iter().zip(config.ae_feature_names.iter()).enumerate() {
|
||||
if mf != sf {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.to_path_buf(),
|
||||
format!("feature[{i}] mismatch: manifest='{mf}' vs sidecar='{sf}'"),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_manifest_overrides(manifest: &ModelManifest, config: &mut MLInferenceConfig) {
|
||||
// Labels: the manifest is authoritative for the label-name map. `confirmations`
|
||||
// and `playbook` are parsed but not consumed here — M2's aggregator rewrite owns that.
|
||||
if !manifest.labels.is_empty() {
|
||||
config.attack_labels = manifest_labels_to_map(&manifest.labels);
|
||||
}
|
||||
// Thresholds: manifest overrides the sidecar when present.
|
||||
if let Some(v) = manifest.thresholds.anomaly {
|
||||
config.anomaly_threshold = v;
|
||||
}
|
||||
if let Some(v) = manifest.thresholds.c2 {
|
||||
config.c2_threshold = v;
|
||||
}
|
||||
if let Some(v) = manifest.thresholds.class_min_confidence {
|
||||
config.class_min_confidence = v;
|
||||
}
|
||||
if let Some(v) = manifest.thresholds.ae {
|
||||
config.ae_threshold = v;
|
||||
}
|
||||
}
|
||||
|
||||
fn manifest_labels_to_map(labels: &std::collections::BTreeMap<String, LabelSpec>) -> HashMap<String, String> {
|
||||
labels.iter().map(|(k, v)| (k.clone(), v.name.clone())).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Integration test: the shipped `models/manifest.yaml` must successfully pair
|
||||
/// with its scaler sidecar to yield a valid `MLInferenceConfig`. Skipped silently
|
||||
/// when run outside the repo root (no `models/` directory).
|
||||
#[test]
|
||||
fn v10_manifest_and_sidecar_load_successfully() {
|
||||
let manifest_path = Path::new("models/manifest.yaml");
|
||||
if !manifest_path.exists() {
|
||||
eprintln!("skipping: models/manifest.yaml absent (not in repo root?)");
|
||||
return;
|
||||
}
|
||||
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(manifest_path)
|
||||
.expect("v10 manifest + sidecar should load cleanly");
|
||||
assert_eq!(manifest.name, "netguardia-v10");
|
||||
assert_eq!(cfg.ae_feature_names.len(), 31);
|
||||
assert_eq!(cfg.classifier_feature_names.len(), 32);
|
||||
// Label map came from manifest, not sidecar.
|
||||
assert_eq!(cfg.attack_labels.get("0").map(String::as_str), Some("Bot"));
|
||||
assert_eq!(cfg.attack_labels.get("7").map(String::as_str), Some("Normal"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feature_mismatch_between_manifest_and_sidecar_is_rejected() {
|
||||
use crate::model::detection::ml_detection::ClipParams;
|
||||
use std::io::Write;
|
||||
|
||||
// Build a minimal sidecar JSON with 2 features.
|
||||
let sidecar = MLInferenceConfig {
|
||||
ae_feature_names: vec!["flow_duration".into(), "fwd_packets".into()],
|
||||
ae_clip_params: HashMap::from([
|
||||
("flow_duration".into(), ClipParams { lower: 0.0, upper: 1.0 }),
|
||||
("fwd_packets".into(), ClipParams { lower: 0.0, upper: 1.0 }),
|
||||
]),
|
||||
ae_scaler_mean: vec![0.0, 0.0],
|
||||
ae_scaler_std: vec![1.0, 1.0],
|
||||
ae_post_clip_min: -5.0,
|
||||
ae_post_clip_max: 5.0,
|
||||
ae_threshold: 0.5,
|
||||
classifier_feature_names: vec!["flow_duration".into(), "fwd_packets".into(), "ae_anomaly_score".into()],
|
||||
attack_labels: HashMap::new(),
|
||||
anomaly_threshold: 0.5,
|
||||
c2_threshold: 0.5,
|
||||
class_min_confidence: 0.4,
|
||||
model_type: "MultiTaskModel".into(),
|
||||
output_names: vec!["anomaly".into(), "class_probs".into(), "c2_score".into()],
|
||||
ae_feature_weights: HashMap::new(),
|
||||
};
|
||||
|
||||
let tmp = std::env::temp_dir().join("netguardia-m1-mismatch-test");
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
let sidecar_path = tmp.join("sidecar.json");
|
||||
let manifest_path = tmp.join("manifest.yaml");
|
||||
let mut f = std::fs::File::create(&sidecar_path).unwrap();
|
||||
f.write_all(serde_json::to_string(&sidecar).unwrap().as_bytes())
|
||||
.unwrap();
|
||||
|
||||
// Manifest lists 3 features, sidecar has 2 — must fail.
|
||||
let manifest_yaml = r#"
|
||||
name: test
|
||||
adapter: multi_task
|
||||
models:
|
||||
autoencoder: ae.onnx
|
||||
classifier: c.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
- fwd_packets
|
||||
- dst_port
|
||||
preprocessing:
|
||||
scaler_sidecar: sidecar.json
|
||||
"#;
|
||||
std::fs::write(&manifest_path, manifest_yaml).unwrap();
|
||||
let err =
|
||||
MLInferenceConfig::from_manifest_with_sidecar(&manifest_path).expect_err("should reject count mismatch");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,9 +1,22 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use common::define::tcp_flags::*;
|
||||
|
||||
use super::flow_tracker::FlowData;
|
||||
use crate::model::detection::flow_features::FlowFeatures;
|
||||
use crate::model::detection::ml_detection::PacketData;
|
||||
|
||||
use crate::model::detection::flow_features::FlowFeatures;
|
||||
/// Signature of a feature getter — takes precomputed flow statistics and returns
|
||||
/// a single f64 feature value. Must be pure (no I/O, no allocation).
|
||||
/// Kept module-private because `PrecomputedStats` is an implementation detail.
|
||||
type FeatureGetter = fn(&PrecomputedStats) -> f64;
|
||||
|
||||
/// Returns true if `name` (canonical or alias) is present in FEATURE_REGISTRY.
|
||||
/// Used by ModelManifest validation at load time.
|
||||
pub fn feature_is_known(name: &str) -> bool {
|
||||
FEATURE_REGISTRY.contains_key(name)
|
||||
}
|
||||
|
||||
impl FlowFeatures {
|
||||
pub fn extract(flow: &FlowData, feature_names: &[String]) -> Self {
|
||||
@ -286,107 +299,206 @@ impl PrecomputedStats {
|
||||
}
|
||||
|
||||
fn get(&self, feature_name: &str) -> f64 {
|
||||
let safe_div = |a: f64, b: f64| if b > 0.0 { a / b } else { 0.0 };
|
||||
|
||||
match feature_name {
|
||||
"Destination Port" | "Dst Port" | "dst_port" => self.dst_port,
|
||||
"Protocol" | "protocol" => self.protocol,
|
||||
"Flow Duration" | "flow_duration" => self.duration_us,
|
||||
"Total Fwd Packets" | "Tot Fwd Pkts" | "fwd_packets" => self.fwd_count,
|
||||
"Total Backward Packets" | "Tot Bwd Pkts" | "bwd_packets" => self.bwd_count,
|
||||
"Total Length of Fwd Packets" | "TotLen Fwd Pkts" | "fwd_bytes" => self.fwd_total_bytes,
|
||||
"Total Length of Bwd Packets" | "TotLen Bwd Pkts" | "bwd_bytes" => self.bwd_total_bytes,
|
||||
"Fwd Packet Length Max" => self.fwd_len_max,
|
||||
"Fwd Packet Length Min" => self.fwd_len_min,
|
||||
"Fwd Packet Length Mean" | "Fwd Pkt Len Mean" | "fwd_pkt_len_mean" => self.fwd_len_mean,
|
||||
"Fwd Packet Length Std" | "Fwd Pkt Len Std" | "fwd_pkt_len_std" => self.fwd_len_std,
|
||||
"Bwd Packet Length Max" => self.bwd_len_max,
|
||||
"Bwd Packet Length Min" => self.bwd_len_min,
|
||||
"Bwd Packet Length Mean" | "Bwd Pkt Len Mean" | "bwd_pkt_len_mean" => self.bwd_len_mean,
|
||||
"Bwd Packet Length Std" | "Bwd Pkt Len Std" | "bwd_pkt_len_std" => self.bwd_len_std,
|
||||
"Flow Bytes/s" | "Flow Byts/s" | "flow_bytes_per_sec" => safe_div(self.total_bytes, self.duration_s),
|
||||
"Flow Packets/s" | "Flow Pkts/s" | "flow_pkts_per_sec" => safe_div(self.total_count, self.duration_s),
|
||||
"Flow IAT Mean" | "flow_iat_mean" => self.flow_iat_mean,
|
||||
"Flow IAT Std" => self.flow_iat_std,
|
||||
"Flow IAT Max" => self.flow_iat_max,
|
||||
"Flow IAT Min" => self.flow_iat_min,
|
||||
"Fwd IAT Total" => self.fwd_iat_total,
|
||||
"Fwd IAT Mean" | "fwd_iat_mean" => self.fwd_iat_mean,
|
||||
"Fwd IAT Std" => self.fwd_iat_std,
|
||||
"Fwd IAT Max" => self.fwd_iat_max,
|
||||
"Fwd IAT Min" => self.fwd_iat_min,
|
||||
"Bwd IAT Total" => self.bwd_iat_total,
|
||||
"Bwd IAT Mean" | "bwd_iat_mean" => self.bwd_iat_mean,
|
||||
"Bwd IAT Std" => self.bwd_iat_std,
|
||||
"Bwd IAT Max" => self.bwd_iat_max,
|
||||
"Bwd IAT Min" => self.bwd_iat_min,
|
||||
"Fwd PSH Flags" => self.fwd_psh,
|
||||
"Bwd PSH Flags" => self.bwd_psh,
|
||||
"Fwd URG Flags" => self.fwd_urg,
|
||||
"Bwd URG Flags" => self.bwd_urg,
|
||||
"Fwd Header Length" => self.fwd_header_bytes,
|
||||
"Bwd Header Length" => self.bwd_header_bytes,
|
||||
"Fwd Packets/s" => safe_div(self.fwd_count, self.duration_s),
|
||||
"Bwd Packets/s" => safe_div(self.bwd_count, self.duration_s),
|
||||
"Min Packet Length" => self.all_len_min,
|
||||
"Max Packet Length" => self.all_len_max,
|
||||
"Packet Length Mean" | "Pkt Len Mean" | "pkt_len_mean" => self.all_len_mean,
|
||||
"Packet Length Std" | "Pkt Len Std" | "pkt_len_std" => self.all_len_std,
|
||||
"Packet Length Variance" => self.all_len_std * self.all_len_std,
|
||||
"FIN Flag Count" | "FIN Flag Cnt" | "fin_flag_cnt" => self.fin_count,
|
||||
"SYN Flag Count" | "SYN Flag Cnt" | "syn_flag_cnt" => self.syn_count,
|
||||
"RST Flag Count" | "RST Flag Cnt" | "rst_flag_cnt" => self.rst_count,
|
||||
"PSH Flag Count" | "PSH Flag Cnt" | "psh_flag_cnt" => self.psh_count,
|
||||
"ACK Flag Count" | "ACK Flag Cnt" | "ack_flag_cnt" => self.ack_count,
|
||||
"URG Flag Count" => self.urg_count,
|
||||
"CWE Flag Count" => self.cwe_count,
|
||||
"ECE Flag Count" => self.ece_count,
|
||||
"Down/Up Ratio" => safe_div(self.bwd_count, self.fwd_count),
|
||||
"Average Packet Size" => safe_div(self.total_bytes, self.total_count),
|
||||
"Avg Fwd Segment Size" => safe_div(self.fwd_total_bytes, self.fwd_count),
|
||||
"Avg Bwd Segment Size" => safe_div(self.bwd_total_bytes, self.bwd_count),
|
||||
"Fwd Header Length.1" => self.fwd_header_bytes,
|
||||
"Fwd Avg Bytes/Bulk" => self.fwd_avg_bytes_bulk,
|
||||
"Fwd Avg Packets/Bulk" => self.fwd_avg_packets_bulk,
|
||||
"Fwd Avg Bulk Rate" => self.fwd_avg_bulk_rate,
|
||||
"Bwd Avg Bytes/Bulk" => self.bwd_avg_bytes_bulk,
|
||||
"Bwd Avg Packets/Bulk" => self.bwd_avg_packets_bulk,
|
||||
"Bwd Avg Bulk Rate" => self.bwd_avg_bulk_rate,
|
||||
"Subflow Fwd Packets" => self.fwd_count,
|
||||
"Subflow Fwd Bytes" => self.fwd_total_bytes,
|
||||
"Subflow Bwd Packets" => self.bwd_count,
|
||||
"Subflow Bwd Bytes" => self.bwd_total_bytes,
|
||||
"fwd_win_bytes" => self.init_win_bytes_fwd,
|
||||
"bwd_win_bytes" => self.init_win_bytes_bwd,
|
||||
"fwd_act_data_pkts" => self.act_data_pkt_fwd,
|
||||
"fwd_seg_size_min" => self.min_seg_size_forward,
|
||||
"Active Mean" => self.active_mean,
|
||||
"Active Std" => self.active_std,
|
||||
"Active Max" => self.active_max,
|
||||
"Active Min" => self.active_min,
|
||||
"Idle Mean" => self.idle_mean,
|
||||
"Idle Std" => self.idle_std,
|
||||
"Idle Max" => self.idle_max,
|
||||
"Idle Min" => self.idle_min,
|
||||
|
||||
// Phase 2: unified names for IAT std (already computed, add aliases)
|
||||
"fwd_iat_std" => self.fwd_iat_std,
|
||||
"bwd_iat_std" => self.bwd_iat_std,
|
||||
"flow_iat_std" => self.flow_iat_std,
|
||||
|
||||
// Phase 2: new features for C2/Bot detection
|
||||
"fwd_bwd_bytes_ratio" => self.fwd_bwd_bytes_ratio,
|
||||
"pkt_len_variance" => self.all_len_std * self.all_len_std,
|
||||
"fwd_iat_skewness" => self.fwd_iat_skewness,
|
||||
|
||||
// Model lists this in ae_feature_names but had no arm — was receiving post-scaler (0-μ)/σ.
|
||||
"iat_cv" => safe_div(self.flow_iat_std, self.flow_iat_mean),
|
||||
|
||||
_ => 0.0,
|
||||
}
|
||||
FEATURE_REGISTRY.get(feature_name).map(|g| g(self)).unwrap_or(0.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn reg_safe_div(a: f64, b: f64) -> f64 {
|
||||
if b > 0.0 { a / b } else { 0.0 }
|
||||
}
|
||||
|
||||
fn reg_insert(m: &mut HashMap<&'static str, FeatureGetter>, names: &[&'static str], g: FeatureGetter) {
|
||||
for n in names {
|
||||
m.insert(*n, g);
|
||||
}
|
||||
}
|
||||
|
||||
/// Central name → getter table. Every name the system recognizes for a feature
|
||||
/// lives here. Manifest validation refuses any name not present in this map.
|
||||
/// Aliases (long-form CICFlowMeter names, short snake_case) map to the same getter.
|
||||
static FEATURE_REGISTRY: LazyLock<HashMap<&'static str, FeatureGetter>> = LazyLock::new(|| {
|
||||
let mut m: HashMap<&'static str, FeatureGetter> = HashMap::new();
|
||||
|
||||
reg_insert(&mut m, &["Destination Port", "Dst Port", "dst_port"], |s| s.dst_port);
|
||||
reg_insert(&mut m, &["Protocol", "protocol"], |s| s.protocol);
|
||||
reg_insert(&mut m, &["Flow Duration", "flow_duration"], |s| s.duration_us);
|
||||
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&[
|
||||
"Total Fwd Packets",
|
||||
"Tot Fwd Pkts",
|
||||
"fwd_packets",
|
||||
"Subflow Fwd Packets",
|
||||
],
|
||||
|s| s.fwd_count,
|
||||
);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&[
|
||||
"Total Backward Packets",
|
||||
"Tot Bwd Pkts",
|
||||
"bwd_packets",
|
||||
"Subflow Bwd Packets",
|
||||
],
|
||||
|s| s.bwd_count,
|
||||
);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&[
|
||||
"Total Length of Fwd Packets",
|
||||
"TotLen Fwd Pkts",
|
||||
"fwd_bytes",
|
||||
"Subflow Fwd Bytes",
|
||||
],
|
||||
|s| s.fwd_total_bytes,
|
||||
);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&[
|
||||
"Total Length of Bwd Packets",
|
||||
"TotLen Bwd Pkts",
|
||||
"bwd_bytes",
|
||||
"Subflow Bwd Bytes",
|
||||
],
|
||||
|s| s.bwd_total_bytes,
|
||||
);
|
||||
|
||||
reg_insert(&mut m, &["Fwd Packet Length Max"], |s| s.fwd_len_max);
|
||||
reg_insert(&mut m, &["Fwd Packet Length Min"], |s| s.fwd_len_min);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&["Fwd Packet Length Mean", "Fwd Pkt Len Mean", "fwd_pkt_len_mean"],
|
||||
|s| s.fwd_len_mean,
|
||||
);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&["Fwd Packet Length Std", "Fwd Pkt Len Std", "fwd_pkt_len_std"],
|
||||
|s| s.fwd_len_std,
|
||||
);
|
||||
|
||||
reg_insert(&mut m, &["Bwd Packet Length Max"], |s| s.bwd_len_max);
|
||||
reg_insert(&mut m, &["Bwd Packet Length Min"], |s| s.bwd_len_min);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&["Bwd Packet Length Mean", "Bwd Pkt Len Mean", "bwd_pkt_len_mean"],
|
||||
|s| s.bwd_len_mean,
|
||||
);
|
||||
reg_insert(
|
||||
&mut m,
|
||||
&["Bwd Packet Length Std", "Bwd Pkt Len Std", "bwd_pkt_len_std"],
|
||||
|s| s.bwd_len_std,
|
||||
);
|
||||
|
||||
reg_insert(&mut m, &["Flow Bytes/s", "Flow Byts/s", "flow_bytes_per_sec"], |s| {
|
||||
reg_safe_div(s.total_bytes, s.duration_s)
|
||||
});
|
||||
reg_insert(&mut m, &["Flow Packets/s", "Flow Pkts/s", "flow_pkts_per_sec"], |s| {
|
||||
reg_safe_div(s.total_count, s.duration_s)
|
||||
});
|
||||
|
||||
reg_insert(&mut m, &["Flow IAT Mean", "flow_iat_mean"], |s| s.flow_iat_mean);
|
||||
reg_insert(&mut m, &["Flow IAT Std", "flow_iat_std"], |s| s.flow_iat_std);
|
||||
reg_insert(&mut m, &["Flow IAT Max"], |s| s.flow_iat_max);
|
||||
reg_insert(&mut m, &["Flow IAT Min"], |s| s.flow_iat_min);
|
||||
|
||||
reg_insert(&mut m, &["Fwd IAT Total"], |s| s.fwd_iat_total);
|
||||
reg_insert(&mut m, &["Fwd IAT Mean", "fwd_iat_mean"], |s| s.fwd_iat_mean);
|
||||
reg_insert(&mut m, &["Fwd IAT Std", "fwd_iat_std"], |s| s.fwd_iat_std);
|
||||
reg_insert(&mut m, &["Fwd IAT Max"], |s| s.fwd_iat_max);
|
||||
reg_insert(&mut m, &["Fwd IAT Min"], |s| s.fwd_iat_min);
|
||||
|
||||
reg_insert(&mut m, &["Bwd IAT Total"], |s| s.bwd_iat_total);
|
||||
reg_insert(&mut m, &["Bwd IAT Mean", "bwd_iat_mean"], |s| s.bwd_iat_mean);
|
||||
reg_insert(&mut m, &["Bwd IAT Std", "bwd_iat_std"], |s| s.bwd_iat_std);
|
||||
reg_insert(&mut m, &["Bwd IAT Max"], |s| s.bwd_iat_max);
|
||||
reg_insert(&mut m, &["Bwd IAT Min"], |s| s.bwd_iat_min);
|
||||
|
||||
reg_insert(&mut m, &["Fwd PSH Flags"], |s| s.fwd_psh);
|
||||
reg_insert(&mut m, &["Bwd PSH Flags"], |s| s.bwd_psh);
|
||||
reg_insert(&mut m, &["Fwd URG Flags"], |s| s.fwd_urg);
|
||||
reg_insert(&mut m, &["Bwd URG Flags"], |s| s.bwd_urg);
|
||||
|
||||
// "Fwd Header Length" and "Fwd Header Length.1" are legacy CICFlowMeter aliases.
|
||||
reg_insert(&mut m, &["Fwd Header Length", "Fwd Header Length.1"], |s| {
|
||||
s.fwd_header_bytes
|
||||
});
|
||||
reg_insert(&mut m, &["Bwd Header Length"], |s| s.bwd_header_bytes);
|
||||
|
||||
reg_insert(&mut m, &["Fwd Packets/s"], |s| reg_safe_div(s.fwd_count, s.duration_s));
|
||||
reg_insert(&mut m, &["Bwd Packets/s"], |s| reg_safe_div(s.bwd_count, s.duration_s));
|
||||
|
||||
reg_insert(&mut m, &["Min Packet Length"], |s| s.all_len_min);
|
||||
reg_insert(&mut m, &["Max Packet Length"], |s| s.all_len_max);
|
||||
reg_insert(&mut m, &["Packet Length Mean", "Pkt Len Mean", "pkt_len_mean"], |s| {
|
||||
s.all_len_mean
|
||||
});
|
||||
reg_insert(&mut m, &["Packet Length Std", "Pkt Len Std", "pkt_len_std"], |s| {
|
||||
s.all_len_std
|
||||
});
|
||||
reg_insert(&mut m, &["Packet Length Variance", "pkt_len_variance"], |s| {
|
||||
s.all_len_std * s.all_len_std
|
||||
});
|
||||
|
||||
reg_insert(&mut m, &["FIN Flag Count", "FIN Flag Cnt", "fin_flag_cnt"], |s| {
|
||||
s.fin_count
|
||||
});
|
||||
reg_insert(&mut m, &["SYN Flag Count", "SYN Flag Cnt", "syn_flag_cnt"], |s| {
|
||||
s.syn_count
|
||||
});
|
||||
reg_insert(&mut m, &["RST Flag Count", "RST Flag Cnt", "rst_flag_cnt"], |s| {
|
||||
s.rst_count
|
||||
});
|
||||
reg_insert(&mut m, &["PSH Flag Count", "PSH Flag Cnt", "psh_flag_cnt"], |s| {
|
||||
s.psh_count
|
||||
});
|
||||
reg_insert(&mut m, &["ACK Flag Count", "ACK Flag Cnt", "ack_flag_cnt"], |s| {
|
||||
s.ack_count
|
||||
});
|
||||
reg_insert(&mut m, &["URG Flag Count"], |s| s.urg_count);
|
||||
reg_insert(&mut m, &["CWE Flag Count"], |s| s.cwe_count);
|
||||
reg_insert(&mut m, &["ECE Flag Count"], |s| s.ece_count);
|
||||
|
||||
reg_insert(&mut m, &["Down/Up Ratio"], |s| reg_safe_div(s.bwd_count, s.fwd_count));
|
||||
reg_insert(&mut m, &["Average Packet Size"], |s| {
|
||||
reg_safe_div(s.total_bytes, s.total_count)
|
||||
});
|
||||
reg_insert(&mut m, &["Avg Fwd Segment Size"], |s| {
|
||||
reg_safe_div(s.fwd_total_bytes, s.fwd_count)
|
||||
});
|
||||
reg_insert(&mut m, &["Avg Bwd Segment Size"], |s| {
|
||||
reg_safe_div(s.bwd_total_bytes, s.bwd_count)
|
||||
});
|
||||
|
||||
reg_insert(&mut m, &["Fwd Avg Bytes/Bulk"], |s| s.fwd_avg_bytes_bulk);
|
||||
reg_insert(&mut m, &["Fwd Avg Packets/Bulk"], |s| s.fwd_avg_packets_bulk);
|
||||
reg_insert(&mut m, &["Fwd Avg Bulk Rate"], |s| s.fwd_avg_bulk_rate);
|
||||
reg_insert(&mut m, &["Bwd Avg Bytes/Bulk"], |s| s.bwd_avg_bytes_bulk);
|
||||
reg_insert(&mut m, &["Bwd Avg Packets/Bulk"], |s| s.bwd_avg_packets_bulk);
|
||||
reg_insert(&mut m, &["Bwd Avg Bulk Rate"], |s| s.bwd_avg_bulk_rate);
|
||||
|
||||
reg_insert(&mut m, &["fwd_win_bytes"], |s| s.init_win_bytes_fwd);
|
||||
reg_insert(&mut m, &["bwd_win_bytes"], |s| s.init_win_bytes_bwd);
|
||||
reg_insert(&mut m, &["fwd_act_data_pkts"], |s| s.act_data_pkt_fwd);
|
||||
reg_insert(&mut m, &["fwd_seg_size_min"], |s| s.min_seg_size_forward);
|
||||
|
||||
reg_insert(&mut m, &["Active Mean"], |s| s.active_mean);
|
||||
reg_insert(&mut m, &["Active Std"], |s| s.active_std);
|
||||
reg_insert(&mut m, &["Active Max"], |s| s.active_max);
|
||||
reg_insert(&mut m, &["Active Min"], |s| s.active_min);
|
||||
reg_insert(&mut m, &["Idle Mean"], |s| s.idle_mean);
|
||||
reg_insert(&mut m, &["Idle Std"], |s| s.idle_std);
|
||||
reg_insert(&mut m, &["Idle Max"], |s| s.idle_max);
|
||||
reg_insert(&mut m, &["Idle Min"], |s| s.idle_min);
|
||||
|
||||
// Phase 2: C2/Bot-oriented features
|
||||
reg_insert(&mut m, &["fwd_bwd_bytes_ratio"], |s| s.fwd_bwd_bytes_ratio);
|
||||
reg_insert(&mut m, &["fwd_iat_skewness"], |s| s.fwd_iat_skewness);
|
||||
reg_insert(&mut m, &["iat_cv"], |s| reg_safe_div(s.flow_iat_std, s.flow_iat_mean));
|
||||
|
||||
m
|
||||
});
|
||||
|
||||
fn compute_stats(values: &[f64]) -> (f64, f64, f64, f64) {
|
||||
if values.is_empty() {
|
||||
return (0.0, 0.0, 0.0, 0.0);
|
||||
@ -469,6 +581,188 @@ mod tests {
|
||||
assert_eq!(compute_bowley_skewness(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), 0.0);
|
||||
}
|
||||
|
||||
/// Hand-crafted PrecomputedStats with distinctive sentinel values per field.
|
||||
/// Lets us verify registry getter dispatch without constructing a real FlowData.
|
||||
fn sample_stats() -> PrecomputedStats {
|
||||
PrecomputedStats {
|
||||
dst_port: 443.0,
|
||||
protocol: 6.0,
|
||||
duration_us: 1_000_000.0,
|
||||
fwd_count: 10.0,
|
||||
bwd_count: 4.0,
|
||||
total_count: 14.0,
|
||||
fwd_total_bytes: 2000.0,
|
||||
bwd_total_bytes: 800.0,
|
||||
total_bytes: 2800.0,
|
||||
duration_s: 1.0,
|
||||
fwd_len_max: 1500.0,
|
||||
fwd_len_min: 40.0,
|
||||
fwd_len_mean: 200.0,
|
||||
fwd_len_std: 300.0,
|
||||
bwd_len_max: 1200.0,
|
||||
bwd_len_min: 60.0,
|
||||
bwd_len_mean: 200.0,
|
||||
bwd_len_std: 250.0,
|
||||
all_len_max: 1500.0,
|
||||
all_len_min: 40.0,
|
||||
all_len_mean: 200.0,
|
||||
all_len_std: 280.0,
|
||||
flow_iat_max: 50_000.0,
|
||||
flow_iat_min: 100.0,
|
||||
flow_iat_mean: 10_000.0,
|
||||
flow_iat_std: 5_000.0,
|
||||
fwd_iat_total: 90_000.0,
|
||||
fwd_iat_max: 40_000.0,
|
||||
fwd_iat_min: 200.0,
|
||||
fwd_iat_mean: 10_000.0,
|
||||
fwd_iat_std: 6_000.0,
|
||||
bwd_iat_total: 30_000.0,
|
||||
bwd_iat_max: 15_000.0,
|
||||
bwd_iat_min: 300.0,
|
||||
bwd_iat_mean: 7_500.0,
|
||||
bwd_iat_std: 4_000.0,
|
||||
fwd_psh: 2.0,
|
||||
bwd_psh: 1.0,
|
||||
fwd_urg: 0.0,
|
||||
bwd_urg: 0.0,
|
||||
fwd_header_bytes: 200.0,
|
||||
bwd_header_bytes: 80.0,
|
||||
fin_count: 1.0,
|
||||
syn_count: 1.0,
|
||||
rst_count: 0.0,
|
||||
psh_count: 3.0,
|
||||
ack_count: 10.0,
|
||||
urg_count: 0.0,
|
||||
cwe_count: 0.0,
|
||||
ece_count: 0.0,
|
||||
fwd_avg_bytes_bulk: 500.0,
|
||||
fwd_avg_packets_bulk: 5.0,
|
||||
fwd_avg_bulk_rate: 5000.0,
|
||||
bwd_avg_bytes_bulk: 400.0,
|
||||
bwd_avg_packets_bulk: 4.0,
|
||||
bwd_avg_bulk_rate: 4000.0,
|
||||
init_win_bytes_fwd: 65535.0,
|
||||
init_win_bytes_bwd: 65000.0,
|
||||
act_data_pkt_fwd: 8.0,
|
||||
min_seg_size_forward: 40.0,
|
||||
active_max: 1000.0,
|
||||
active_min: 50.0,
|
||||
active_mean: 300.0,
|
||||
active_std: 200.0,
|
||||
idle_max: 500.0,
|
||||
idle_min: 10.0,
|
||||
idle_mean: 100.0,
|
||||
idle_std: 80.0,
|
||||
fwd_bwd_bytes_ratio: 0.71,
|
||||
fwd_iat_skewness: 0.15,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_unknown_name_returns_zero() {
|
||||
let s = sample_stats();
|
||||
assert_eq!(s.get("not_a_feature"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_aliases_resolve_identically() {
|
||||
// Long-form, short-form, and "Subflow" aliases must all map to the same getter.
|
||||
let s = sample_stats();
|
||||
for group in [
|
||||
[
|
||||
"Total Fwd Packets",
|
||||
"Tot Fwd Pkts",
|
||||
"fwd_packets",
|
||||
"Subflow Fwd Packets",
|
||||
],
|
||||
[
|
||||
"Total Length of Fwd Packets",
|
||||
"TotLen Fwd Pkts",
|
||||
"fwd_bytes",
|
||||
"Subflow Fwd Bytes",
|
||||
],
|
||||
["Flow IAT Std", "flow_iat_std", "Flow IAT Std", "Flow IAT Std"], // pad to 4
|
||||
["Fwd IAT Std", "fwd_iat_std", "Fwd IAT Std", "Fwd IAT Std"],
|
||||
["Bwd IAT Std", "bwd_iat_std", "Bwd IAT Std", "Bwd IAT Std"],
|
||||
[
|
||||
"Packet Length Variance",
|
||||
"pkt_len_variance",
|
||||
"Packet Length Variance",
|
||||
"Packet Length Variance",
|
||||
],
|
||||
[
|
||||
"Fwd Header Length",
|
||||
"Fwd Header Length.1",
|
||||
"Fwd Header Length",
|
||||
"Fwd Header Length",
|
||||
],
|
||||
] {
|
||||
let expected = s.get(group[0]);
|
||||
for name in &group[1..] {
|
||||
assert_eq!(
|
||||
s.get(name),
|
||||
expected,
|
||||
"alias '{name}' should resolve to same value as '{}'",
|
||||
group[0]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_safe_div_returns_zero_on_zero_denominator() {
|
||||
let mut s = sample_stats();
|
||||
s.flow_iat_mean = 0.0;
|
||||
s.flow_iat_std = 500.0;
|
||||
// iat_cv = std / mean, but mean=0 → safe_div → 0.0
|
||||
assert_eq!(s.get("iat_cv"), 0.0);
|
||||
s.duration_s = 0.0;
|
||||
assert_eq!(s.get("flow_bytes_per_sec"), 0.0);
|
||||
assert_eq!(s.get("flow_pkts_per_sec"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registry_covers_v10_manifest_features() {
|
||||
// Every feature the shipped v10 manifest references must be registered.
|
||||
// A missing name here means the match → registry refactor dropped a binding.
|
||||
const V10_FEATURES: &[&str] = &[
|
||||
"flow_duration",
|
||||
"fwd_packets",
|
||||
"bwd_packets",
|
||||
"fwd_bytes",
|
||||
"bwd_bytes",
|
||||
"flow_bytes_per_sec",
|
||||
"flow_pkts_per_sec",
|
||||
"fwd_win_bytes",
|
||||
"bwd_win_bytes",
|
||||
"fwd_pkt_len_mean",
|
||||
"bwd_pkt_len_mean",
|
||||
"fwd_iat_mean",
|
||||
"bwd_iat_mean",
|
||||
"flow_iat_mean",
|
||||
"pkt_len_mean",
|
||||
"dst_port",
|
||||
"protocol",
|
||||
"psh_flag_cnt",
|
||||
"ack_flag_cnt",
|
||||
"syn_flag_cnt",
|
||||
"fin_flag_cnt",
|
||||
"rst_flag_cnt",
|
||||
"pkt_len_std",
|
||||
"fwd_pkt_len_std",
|
||||
"bwd_pkt_len_std",
|
||||
"fwd_seg_size_min",
|
||||
"fwd_act_data_pkts",
|
||||
"fwd_iat_std",
|
||||
"bwd_iat_std",
|
||||
"fwd_bwd_bytes_ratio",
|
||||
"iat_cv",
|
||||
];
|
||||
for f in V10_FEATURES {
|
||||
assert!(feature_is_known(f), "v10 feature '{f}' missing from FEATURE_REGISTRY");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bowley_skewness_known_output() {
|
||||
// Symmetric distribution: [1, 2, 3, 4, 5, 6, 7, 8] (n=8)
|
||||
|
||||
229
net-guardia/src/core/ml/manifest.rs
Normal file
229
net-guardia/src/core/ml/manifest.rs
Normal file
@ -0,0 +1,229 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::feature_extractor::feature_is_known;
|
||||
use crate::model::error::ml::MLError;
|
||||
|
||||
/// The three fixed adapter shapes a v1 manifest may declare.
|
||||
/// The user picks one string; never sees the internal `ModelAdapter` trait.
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum AdapterKind {
|
||||
ClassifierOnly,
|
||||
AutoencoderOnly,
|
||||
MultiTask,
|
||||
}
|
||||
|
||||
impl AdapterKind {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
AdapterKind::ClassifierOnly => "classifier_only",
|
||||
AdapterKind::AutoencoderOnly => "autoencoder_only",
|
||||
AdapterKind::MultiTask => "multi_task",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ONNX file references within a manifest. `classifier_only` / `autoencoder_only`
|
||||
/// adapters use `model`; `multi_task` uses `autoencoder` + `classifier`.
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
pub struct ModelPaths {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub autoencoder: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub classifier: Option<String>,
|
||||
}
|
||||
|
||||
/// Per-label metadata. `confirmations` and `playbook` are parsed but not yet consumed
|
||||
/// by the aggregator — they belong to M2's aggregator rewrite.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct LabelSpec {
|
||||
pub name: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub confirmations: Option<usize>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub playbook: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
pub struct Thresholds {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub anomaly: Option<f32>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub c2: Option<f32>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub class_min_confidence: Option<f32>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub ae: Option<f32>,
|
||||
}
|
||||
|
||||
/// Pointer to a preprocessing sidecar (scaler/clip arrays). For v10 this is
|
||||
/// the legacy `inference_config.json`.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Preprocessing {
|
||||
pub scaler_sidecar: String,
|
||||
}
|
||||
|
||||
/// The user-authored YAML description of a model. Validated at load — unknown
|
||||
/// feature names are rejected against `FEATURE_REGISTRY`.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct ModelManifest {
|
||||
pub name: String,
|
||||
pub adapter: AdapterKind,
|
||||
#[serde(default)]
|
||||
pub models: ModelPaths,
|
||||
pub features: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub labels: BTreeMap<String, LabelSpec>,
|
||||
#[serde(default)]
|
||||
pub thresholds: Thresholds,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub preprocessing: Option<Preprocessing>,
|
||||
}
|
||||
|
||||
impl ModelManifest {
|
||||
/// Parse and validate a manifest YAML file from disk.
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self, MLError> {
|
||||
let path = path.as_ref();
|
||||
let content = fs::read_to_string(path)
|
||||
.map_err(|e| MLError::ManifestInvalid(path.to_path_buf(), format!("read failed: {e}")))?;
|
||||
let manifest: ModelManifest = serde_yaml_ng::from_str(&content)
|
||||
.map_err(|e| MLError::ManifestInvalid(path.to_path_buf(), format!("YAML parse: {e}")))?;
|
||||
manifest.validate(path)?;
|
||||
Ok(manifest)
|
||||
}
|
||||
|
||||
fn validate(&self, path: &Path) -> Result<(), MLError> {
|
||||
if self.name.trim().is_empty() {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
"name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.features.is_empty() {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
"features is empty".to_string(),
|
||||
));
|
||||
}
|
||||
for f in &self.features {
|
||||
if !feature_is_known(f) {
|
||||
return Err(MLError::UnknownFeature(f.clone()));
|
||||
}
|
||||
}
|
||||
match self.adapter {
|
||||
AdapterKind::MultiTask => {
|
||||
if self.models.autoencoder.is_none() || self.models.classifier.is_none() {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
"multi_task adapter requires both models.autoencoder and models.classifier".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
AdapterKind::ClassifierOnly | AdapterKind::AutoencoderOnly => {
|
||||
if self.models.model.is_none() {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
format!("{} adapter requires models.model", self.adapter.as_str()),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve a manifest-relative filename (e.g. "deep_autoencoder.onnx") to an
|
||||
/// absolute-ish path rooted at the manifest's parent directory.
|
||||
pub fn resolve_relative(manifest_path: &Path, relative: &str) -> PathBuf {
|
||||
manifest_path.parent().unwrap_or_else(|| Path::new(".")).join(relative)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const V10_MANIFEST: &str = r#"
|
||||
name: netguardia-v10
|
||||
adapter: multi_task
|
||||
models:
|
||||
autoencoder: deep_autoencoder.onnx
|
||||
classifier: classifier.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
- fwd_packets
|
||||
- bwd_packets
|
||||
labels:
|
||||
"0": { name: Bot }
|
||||
"7": { name: Normal }
|
||||
thresholds:
|
||||
anomaly: 0.9
|
||||
c2: 0.85
|
||||
preprocessing:
|
||||
scaler_sidecar: inference_config.json
|
||||
"#;
|
||||
|
||||
#[test]
|
||||
fn parses_minimal_multitask() {
|
||||
let m: ModelManifest = serde_yaml_ng::from_str(V10_MANIFEST).expect("parse");
|
||||
assert_eq!(m.name, "netguardia-v10");
|
||||
assert_eq!(m.adapter, AdapterKind::MultiTask);
|
||||
assert_eq!(m.features.len(), 3);
|
||||
assert_eq!(m.models.autoencoder.as_deref(), Some("deep_autoencoder.onnx"));
|
||||
assert_eq!(m.models.classifier.as_deref(), Some("classifier.onnx"));
|
||||
assert_eq!(m.labels.len(), 2);
|
||||
assert_eq!(m.labels.get("0").map(|l| l.name.as_str()), Some("Bot"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unknown_feature() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: classifier_only
|
||||
models:
|
||||
model: m.onnx
|
||||
features:
|
||||
- this_feature_does_not_exist
|
||||
"#;
|
||||
let path = Path::new("/tmp/test-manifest.yaml");
|
||||
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
|
||||
let err = parsed.validate(path).expect_err("should reject unknown feature");
|
||||
assert!(matches!(err, MLError::UnknownFeature { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_multitask_missing_ae() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: multi_task
|
||||
models:
|
||||
classifier: c.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
"#;
|
||||
let path = Path::new("/tmp/test-manifest.yaml");
|
||||
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
|
||||
let err = parsed.validate(path).expect_err("should require autoencoder");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_features() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: classifier_only
|
||||
models:
|
||||
model: m.onnx
|
||||
features: []
|
||||
"#;
|
||||
let path = Path::new("/tmp/test-manifest.yaml");
|
||||
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
|
||||
let err = parsed.validate(path).expect_err("should reject empty features");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
}
|
||||
@ -6,6 +6,7 @@ pub mod engine;
|
||||
pub mod feature_extractor;
|
||||
pub mod flow_tracker;
|
||||
pub mod inference;
|
||||
pub mod manifest;
|
||||
pub mod model_loader;
|
||||
pub mod model_watcher;
|
||||
pub mod traffic_logger;
|
||||
|
||||
@ -3,6 +3,8 @@ use std::time::Instant;
|
||||
|
||||
use macros::log;
|
||||
use tract_onnx::prelude::*;
|
||||
use tract_onnx::tract_hir::infer::Factoid;
|
||||
use tract_onnx::tract_hir::internal::DimLike;
|
||||
|
||||
use crate::infrastructure::app_config::AppConfig;
|
||||
use crate::model::detection::ml_detection::RunnableModel;
|
||||
@ -21,18 +23,26 @@ impl MLModels {
|
||||
app_config: &Arc<AppConfig>,
|
||||
inference_config: &Arc<MLInferenceConfig>,
|
||||
batch_size: usize,
|
||||
) -> Result<Self, MLError> {
|
||||
Self::load_named(
|
||||
&app_config.inference.deep_autoencoder_name,
|
||||
&app_config.inference.classifier_name,
|
||||
inference_config,
|
||||
batch_size,
|
||||
)
|
||||
}
|
||||
|
||||
/// Manifest-friendly constructor: caller supplies the AE and classifier ONNX filenames
|
||||
/// directly. Used when a `ModelManifest` overrides the defaults in `AppConfig`.
|
||||
pub fn load_named(
|
||||
autoencoder_name: &str,
|
||||
classifier_name: &str,
|
||||
inference_config: &Arc<MLInferenceConfig>,
|
||||
batch_size: usize,
|
||||
) -> Result<Self, MLError> {
|
||||
Ok(Self {
|
||||
deep_autoencoder: Self::loader(
|
||||
&app_config.inference.deep_autoencoder_name,
|
||||
inference_config.num_ae_features(),
|
||||
batch_size,
|
||||
)?,
|
||||
classifier: Self::loader(
|
||||
&app_config.inference.classifier_name,
|
||||
inference_config.num_classifier_features(),
|
||||
batch_size,
|
||||
)?,
|
||||
deep_autoencoder: Self::loader(autoencoder_name, inference_config.num_ae_features(), batch_size)?,
|
||||
classifier: Self::loader(classifier_name, inference_config.num_classifier_features(), batch_size)?,
|
||||
batch_size,
|
||||
})
|
||||
}
|
||||
@ -43,21 +53,44 @@ impl MLModels {
|
||||
log!(MLLog::ModelLoading(model.to_string(), features, batch_size));
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
let load = || -> Result<RunnableModel, Box<dyn std::error::Error>> {
|
||||
let mut model = onnx().model_for_path(&model_path)?;
|
||||
model.set_input_fact(0, f32::fact([batch_size, features]).into())?;
|
||||
Ok(model.into_optimized()?.into_runnable()?)
|
||||
};
|
||||
|
||||
let result = load().map_err(|_| MLError::ModelLoadFailed(model_path));
|
||||
|
||||
let result = Self::loader_inner(&model_path, model, features, batch_size);
|
||||
let elapsed_ms = start.elapsed().as_millis() as u64;
|
||||
log!(MLLog::ModelLoadComplete(model.to_string(), elapsed_ms));
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn loader_inner(
|
||||
model_path: &PathBuf,
|
||||
model: &str,
|
||||
features: usize,
|
||||
batch_size: usize,
|
||||
) -> Result<RunnableModel, MLError> {
|
||||
let mut onnx_model = onnx()
|
||||
.model_for_path(model_path)
|
||||
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))?;
|
||||
|
||||
// Introspect the ONNX input fact and cross-check its last dim against the
|
||||
// expected feature count. A dynamic/symbolic dim is skipped — we only fail
|
||||
// when the ONNX graph declares a concrete integer that disagrees.
|
||||
if let Some(onnx_dim) = introspect_input_features(&onnx_model) {
|
||||
let matched = onnx_dim == features;
|
||||
log!(MLLog::OnnxShapeChecked(model.to_string(), features, onnx_dim, matched,));
|
||||
if !matched {
|
||||
return Err(MLError::FeatureMismatch(model_path.clone(), features, onnx_dim));
|
||||
}
|
||||
}
|
||||
|
||||
onnx_model
|
||||
.set_input_fact(0, f32::fact([batch_size, features]).into())
|
||||
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))?;
|
||||
|
||||
onnx_model
|
||||
.into_optimized()
|
||||
.and_then(|m| m.into_runnable())
|
||||
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))
|
||||
}
|
||||
|
||||
pub fn get_model_info(&self, name: &str) -> String {
|
||||
let model = match name {
|
||||
"deep_autoencoder" => &self.deep_autoencoder,
|
||||
@ -73,3 +106,67 @@ impl MLModels {
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the concrete last-dim (feature count) from an ONNX model's declared input fact.
|
||||
/// Returns None when the dim is dynamic/symbolic or when the model has no input 0.
|
||||
fn introspect_input_features(model: &InferenceModel) -> Option<usize> {
|
||||
let fact = model.input_fact(0).ok()?;
|
||||
let rank = fact.shape.rank().concretize()? as usize;
|
||||
if rank == 0 {
|
||||
return None;
|
||||
}
|
||||
let last = fact.shape.dim(rank - 1)?;
|
||||
last.concretize()?.to_usize().ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Integration test: tract must introspect the shipped v10 AE ONNX's input dim as 31.
|
||||
/// Silently skipped outside the repo root (no `models/` directory).
|
||||
#[test]
|
||||
fn introspect_v10_autoencoder_input_is_31() {
|
||||
let ae_path = PathBuf::from("models/deep_autoencoder.onnx");
|
||||
if !ae_path.exists() {
|
||||
eprintln!("skipping: models/deep_autoencoder.onnx absent");
|
||||
return;
|
||||
}
|
||||
let model = onnx().model_for_path(&ae_path).expect("load AE onnx");
|
||||
let dim = introspect_input_features(&model).expect("v10 AE should expose a concrete final-dim");
|
||||
assert_eq!(dim, 31, "v10 AE ONNX input dim changed unexpectedly");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn introspect_v10_classifier_input_is_32() {
|
||||
let cls_path = PathBuf::from("models/classifier.onnx");
|
||||
if !cls_path.exists() {
|
||||
eprintln!("skipping: models/classifier.onnx absent");
|
||||
return;
|
||||
}
|
||||
let model = onnx().model_for_path(&cls_path).expect("load classifier onnx");
|
||||
let dim = introspect_input_features(&model).expect("v10 classifier should expose a concrete final-dim");
|
||||
assert_eq!(dim, 32, "v10 classifier ONNX input dim changed unexpectedly");
|
||||
}
|
||||
|
||||
/// End-to-end M1 smoke: loading v10 via `load_named` with the manifest-driven
|
||||
/// names + sidecar-derived inference config must succeed without shape mismatch.
|
||||
#[test]
|
||||
fn v10_load_named_with_manifest_paths_succeeds() {
|
||||
let manifest_path = std::path::Path::new("models/manifest.yaml");
|
||||
if !manifest_path.exists() {
|
||||
eprintln!("skipping: models/manifest.yaml absent");
|
||||
return;
|
||||
}
|
||||
let (cfg, manifest) =
|
||||
MLInferenceConfig::from_manifest_with_sidecar(manifest_path).expect("manifest + sidecar load");
|
||||
let cfg = Arc::new(cfg);
|
||||
let ae = manifest.models.autoencoder.as_deref().expect("multi_task has AE");
|
||||
let cls = manifest
|
||||
.models
|
||||
.classifier
|
||||
.as_deref()
|
||||
.expect("multi_task has classifier");
|
||||
let _models = MLModels::load_named(ae, cls, &cfg, 8).expect("v10 load via manifest names");
|
||||
}
|
||||
}
|
||||
|
||||
@ -359,11 +359,8 @@ impl System {
|
||||
// Start Suricata eve.json monitor — tails the log file, translates
|
||||
// alert events into DetectionEvent on the shared mpsc. No-op if the
|
||||
// bridge is disabled in config.
|
||||
crate::infrastructure::suricata_monitor::SuricataMonitor::new(
|
||||
self.app_config.clone(),
|
||||
suricata_detection_tx,
|
||||
)
|
||||
.start();
|
||||
crate::infrastructure::suricata_monitor::SuricataMonitor::new(self.app_config.clone(), suricata_detection_tx)
|
||||
.start();
|
||||
|
||||
// Wait for shutdown signal (ctrl-c OR API-triggered)
|
||||
tokio::select! {
|
||||
|
||||
@ -8,6 +8,7 @@ use tokio::sync::oneshot;
|
||||
use crate::core::ml::alert::MLAlert;
|
||||
use crate::core::ml::drift_detector::DriftDetector;
|
||||
use crate::core::ml::engine::Engine;
|
||||
use crate::core::ml::manifest::{AdapterKind, ModelManifest};
|
||||
use crate::core::ml::model_loader::MLModels;
|
||||
use crate::core::ml::traffic_logger::TrafficLogger;
|
||||
use crate::infrastructure::app_config::AppConfig;
|
||||
@ -38,13 +39,32 @@ impl AppServices {
|
||||
pub fn new(
|
||||
app_config: Arc<AppConfig>,
|
||||
inference_config: Arc<MLInferenceConfig>,
|
||||
ml_manifest: Option<ModelManifest>,
|
||||
drift_detector: Arc<parking_lot::Mutex<DriftDetector>>,
|
||||
ebpf_health: Arc<parking_lot::RwLock<EbpfHealth>>,
|
||||
) -> Result<Self, Error> {
|
||||
let health = SystemHealth::new(app_config.clone(), ebpf_health)?;
|
||||
|
||||
let batch_size = app_config.inference.inference_batch_size;
|
||||
let ml_models = Arc::new(MLModels::load_models(&app_config, &inference_config, batch_size)?);
|
||||
// When a `multi_task` manifest is present, its model filenames are authoritative.
|
||||
// Otherwise (no manifest, or single-model adapter in v2) fall through to the
|
||||
// legacy app_config-driven path. M1 only needs `multi_task` wired end-to-end.
|
||||
let ml_models = Arc::new(match ml_manifest.as_ref() {
|
||||
Some(m) if m.adapter == AdapterKind::MultiTask => {
|
||||
let ae = m
|
||||
.models
|
||||
.autoencoder
|
||||
.as_deref()
|
||||
.unwrap_or(&app_config.inference.deep_autoencoder_name);
|
||||
let cls = m
|
||||
.models
|
||||
.classifier
|
||||
.as_deref()
|
||||
.unwrap_or(&app_config.inference.classifier_name);
|
||||
MLModels::load_named(ae, cls, &inference_config, batch_size)?
|
||||
}
|
||||
_ => MLModels::load_models(&app_config, &inference_config, batch_size)?,
|
||||
});
|
||||
let ml_alert = Arc::new(MLAlert::new());
|
||||
|
||||
let traffic_logger = if app_config.inference.traffic_logging_mode {
|
||||
|
||||
@ -155,8 +155,14 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn categorizes_permission_errors() {
|
||||
assert!(matches!(categorize("Permission denied (os error 13)"), EbpfFailCategory::Permission));
|
||||
assert!(matches!(categorize("Operation not permitted"), EbpfFailCategory::Permission));
|
||||
assert!(matches!(
|
||||
categorize("Permission denied (os error 13)"),
|
||||
EbpfFailCategory::Permission
|
||||
));
|
||||
assert!(matches!(
|
||||
categorize("Operation not permitted"),
|
||||
EbpfFailCategory::Permission
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@ -147,7 +147,20 @@ impl SecretStore {
|
||||
.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SecretStorePort for SecretStore {
|
||||
fn get_secret(&self, key: &str) -> Result<Option<String>, Error> {
|
||||
match self.db.get_app_secret(key)? {
|
||||
Some(envelope_json) => Ok(Some(self.decrypt(&envelope_json)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_secret(&self, key: &str, plaintext: &str) -> Result<(), Error> {
|
||||
let envelope = self.encrypt(plaintext)?;
|
||||
self.db.set_app_secret(key, &envelope)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@ -255,17 +268,3 @@ mod tests {
|
||||
assert!(store.decrypt(&envelope).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
impl SecretStorePort for SecretStore {
|
||||
fn get_secret(&self, key: &str) -> Result<Option<String>, Error> {
|
||||
match self.db.get_app_secret(key)? {
|
||||
Some(envelope_json) => Ok(Some(self.decrypt(&envelope_json)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_secret(&self, key: &str, plaintext: &str) -> Result<(), Error> {
|
||||
let envelope = self.encrypt(plaintext)?;
|
||||
self.db.set_app_secret(key, &envelope)
|
||||
}
|
||||
}
|
||||
|
||||
@ -20,6 +20,7 @@ use crate::core::dns_filter_service::DnsFilterService;
|
||||
use crate::core::ebpf::EbpfServices;
|
||||
use crate::core::email::scheduler::ReportScheduler;
|
||||
use crate::core::ml::drift_detector::DriftDetector;
|
||||
use crate::core::ml::manifest::ModelManifest;
|
||||
use crate::core::notification_service::NotificationService;
|
||||
use crate::core::playbook_service::PlaybookService;
|
||||
use crate::core::rate_limit_service::RateLimitService;
|
||||
@ -44,6 +45,7 @@ use crate::model::error::Error;
|
||||
use crate::model::error::ebpf::EbpfError;
|
||||
use crate::model::error::misc::MiscError;
|
||||
use crate::model::log::ebpf::EbpfLog;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::log::system::SystemLog;
|
||||
use crate::model::monitoring::direction::FlowDirection;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
@ -104,7 +106,26 @@ impl ServiceFactory {
|
||||
AppConfig::seed_defaults(&db)?;
|
||||
let app_config = Arc::new(AppConfig::new(&db)?);
|
||||
|
||||
let inference_config = Arc::new(MLInferenceConfig::load_file(&app_config.inference.models_config_name)?);
|
||||
// Prefer `models/manifest.yaml` when present (v12 BYO-model path). The manifest
|
||||
// is the user-authored source of truth for features, labels, thresholds, and
|
||||
// model filenames; the legacy JSON-only path is the fallback.
|
||||
let manifest_path = std::path::PathBuf::from("models/manifest.yaml");
|
||||
let (inference_config, ml_manifest): (Arc<MLInferenceConfig>, Option<ModelManifest>) = if manifest_path.exists()
|
||||
{
|
||||
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(&manifest_path)?;
|
||||
log!(MLLog::ManifestLoaded(
|
||||
manifest.name.clone(),
|
||||
manifest.adapter.as_str().to_string(),
|
||||
manifest.features.len(),
|
||||
manifest.labels.len(),
|
||||
));
|
||||
(Arc::new(cfg), Some(manifest))
|
||||
} else {
|
||||
(
|
||||
Arc::new(MLInferenceConfig::load_file(&app_config.inference.models_config_name)?),
|
||||
None,
|
||||
)
|
||||
};
|
||||
|
||||
// Shared eBPF health handle. Initialized Healthy; downgraded to
|
||||
// Unavailable with a classified reason if any stage below fails.
|
||||
@ -116,16 +137,21 @@ impl ServiceFactory {
|
||||
// rest of the system (HTTP API, SOAR, ML engine, auth) is built
|
||||
// regardless so the operator can still reach the frontend and see
|
||||
// the reason.
|
||||
let (ingress_ebpf, egress_ebpf, ingress_program_array, ebpf_services) =
|
||||
match Self::try_build_ebpf(&app_config) {
|
||||
Ok((ingress, egress, pa, services)) => (Some(ingress), Some(egress), Some(pa), Arc::new(services)),
|
||||
Err((stage, err)) => {
|
||||
let health = crate::infrastructure::ebpf_preflight::classify(stage, &err, None);
|
||||
log!(SystemLog::EbpfBringupFailed(format!("{:?}", health)));
|
||||
*ebpf_health.write() = health;
|
||||
(None, None, None, Arc::new(EbpfServices::unavailable(app_config.clone())))
|
||||
}
|
||||
};
|
||||
let (ingress_ebpf, egress_ebpf, ingress_program_array, ebpf_services) = match Self::try_build_ebpf(&app_config)
|
||||
{
|
||||
Ok((ingress, egress, pa, services)) => (Some(ingress), Some(egress), Some(pa), Arc::new(services)),
|
||||
Err((stage, err)) => {
|
||||
let health = crate::infrastructure::ebpf_preflight::classify(stage, &err, None);
|
||||
log!(SystemLog::EbpfBringupFailed(format!("{:?}", health)));
|
||||
*ebpf_health.write() = health;
|
||||
(
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Arc::new(EbpfServices::unavailable(app_config.clone())),
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
// Ensure enforce_mode setting exists (default: monitor)
|
||||
if db.get_setting("enforce_mode")?.is_none() {
|
||||
@ -153,6 +179,7 @@ impl ServiceFactory {
|
||||
let app_services = Arc::new(AppServices::new(
|
||||
app_config.clone(),
|
||||
inference_config.clone(),
|
||||
ml_manifest.clone(),
|
||||
drift_detector.clone(),
|
||||
ebpf_health.clone(),
|
||||
)?);
|
||||
@ -267,8 +294,7 @@ impl ServiceFactory {
|
||||
secret_store_port,
|
||||
));
|
||||
|
||||
let suricata_manager =
|
||||
crate::infrastructure::suricata_manager::SuricataManager::new(app_config.clone());
|
||||
let suricata_manager = crate::infrastructure::suricata_manager::SuricataManager::new(app_config.clone());
|
||||
|
||||
Ok(AppState {
|
||||
app_config,
|
||||
|
||||
@ -71,18 +71,14 @@ impl SuricataManager {
|
||||
loop {
|
||||
// Pre-flight: validate binary + config exist before spawning.
|
||||
if let Err(e) = Self::preflight(&self.config) {
|
||||
*self.health.write() = SuricataHealth::Stopped {
|
||||
reason: e.to_string(),
|
||||
};
|
||||
*self.health.write() = SuricataHealth::Stopped { reason: e.to_string() };
|
||||
return;
|
||||
}
|
||||
|
||||
let mut child = match self.spawn_child() {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
*self.health.write() = SuricataHealth::Stopped {
|
||||
reason: e.to_string(),
|
||||
};
|
||||
*self.health.write() = SuricataHealth::Stopped { reason: e.to_string() };
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
@ -144,7 +144,11 @@ impl SuricataMonitor {
|
||||
|
||||
let alert = v.get("alert")?;
|
||||
let sid = alert.get("signature_id").and_then(|x| x.as_u64()).unwrap_or(0) as u32;
|
||||
let signature = alert.get("signature").and_then(|x| x.as_str()).unwrap_or("").to_string();
|
||||
let signature = alert
|
||||
.get("signature")
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
// Suricata severity: 1=high, 2=medium, 3=low, 4=informational.
|
||||
// Map to confidence in [0.5, 1.0] so high-severity alerts tend to trip
|
||||
// SOAR condition thresholds.
|
||||
|
||||
@ -22,5 +22,17 @@ traceable! {
|
||||
|
||||
#[error("Failed to flush traffic log: {err}")]
|
||||
TrafficLogFlushFailed => tracing::Level::ERROR,
|
||||
|
||||
#[no_source]
|
||||
#[error("Model manifest at {path:?} is invalid: {reason}")]
|
||||
ManifestInvalid { path: PathBuf, reason: String } => tracing::Level::ERROR,
|
||||
|
||||
#[no_source]
|
||||
#[error("Feature count mismatch for {model:?}: manifest declares {declared}, ONNX input expects {onnx_dim}")]
|
||||
FeatureMismatch { model: PathBuf, declared: usize, onnx_dim: usize } => tracing::Level::ERROR,
|
||||
|
||||
#[no_source]
|
||||
#[error("Unknown feature '{name}' — not registered in FEATURE_REGISTRY")]
|
||||
UnknownFeature { name: String } => tracing::Level::ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
@ -62,5 +62,11 @@ loggable! {
|
||||
|
||||
#[error("Model reload failed, keeping current models: {error}")]
|
||||
ModelReloadFailed { error: String } => tracing::Level::ERROR,
|
||||
|
||||
#[error("Model manifest loaded: name='{name}', adapter={adapter}, features={features}, labels={labels}")]
|
||||
ManifestLoaded { name: String, adapter: String, features: usize, labels: usize } => tracing::Level::INFO,
|
||||
|
||||
#[error("ONNX input shape introspected for {model}: declared={declared}, onnx_dim={onnx_dim}, matched={matched}")]
|
||||
OnnxShapeChecked { model: String, declared: usize, onnx_dim: usize, matched: bool } => tracing::Level::DEBUG,
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user