mirror of
https://github.com/DaLaw2/NetGuardia.git
synced 2026-08-24 14:10:28 +09:00
feat(core): BYO ML adapter state machine + Fusion v1 foundation
Cross-cutting rework that lands the multi-source NIDS+IPS fusion groundwork
alongside a manifest-driven ML adapter. Day 1 with no model loaded is a
first-class state, and the remaining 3 sources (Suricata + CV + Graph) fuse
on a shared vocabulary.
ML pipeline (BYO-ready)
- Replace the hardcoded `MLModels { deep_autoencoder, classifier, batch_size }`
struct with `MLModelAdapter` enum (AutoencoderOnly / ClassifierOnly / MultiTask),
wrapped in `ModelSourceState { Dormant | Active { adapter, info } | Error { msg,
since, last_attempted_path } }` and held inside `ArcSwap` on `Inference`.
Day 1 without a manifest resolves to `Dormant`; reload failures park in
`Error` carrying the reason and offending path.
- `MultiTask` arm now reads the c2 label index from the manifest instead of
matching the hardcoded "C2 Communication" string, and replaces the
`unwrap_or(0.0)` sentinel with an explicit `if let && let && c2 > c2_class_prob`
chain.
- Aggregator `should_alert` takes `required_confirmations: usize` and
`alert_multiplier: f32` directly from manifest `labels[*].confirmations` /
`thresholds.alert_multiplier`; the per-attack-type string match and
`ALERT_THRESHOLD_MULTIPLIER` constant are gone.
- Manifest validation rejects `confirmations: 0`, case-insensitive duplicate
label names, and `alert_multiplier <= 0`.
- New `ModelSourceStatus` wire type (serde-tagged Dormant/Active/Error) for
HTTP / WebSocket consumers.
Fusion core
- Canonical attack-type dictionary (13 seed types + `Unknown` bucket) with
per-source `translate(source, raw_label)` and `canonical_from_str`. ML /
Beaconing / Correlation / Suricata raw labels fold into a shared vocabulary
so the orchestrator's `(source_ip, canonical_attack_type)` dedup key
actually collides across sources.
- `fusion_math::fused_confidence(&[f32]) -> f32` implements `1 - prod(1 - c_i)`,
clamped to `[0.0, 1.0]` and safe for empty input.
- `FusionWindowLengths` pins dynamic per-first-source window lengths
(Suricata 10s / CV 8s / ML 2s / Graph 5s) with `[1, 30]` second clamp.
- Orchestrator canonicalizes raw labels on ingress, opens a fusion window
keyed by (src_ip, canonical), and silences duplicates for 30s after emit.
- SOAR `Condition::MultiSourceMin { min }` and
`Condition::SingleSourceHigh { source, min_confidence }` implement the
"N sources agreed" gate and solo-high-confidence escape hatch.
`SoarEngine::reload_cache` canonicalizes `trigger_event` on load and logs
`SoarLog::NonCanonicalTriggerEvent` for anything outside the vocabulary.
Workspace hygiene
- `MODELS_DIR` / `MANIFEST_FILENAME` / `STAGING_SUBDIR` hoisted to
`model::config::constants`.
- `DetectionSource` gains `FromStr`; `soar/matcher.rs` string matches on
source become parser-driven.
- Magic numbers in `soar/matcher.rs` (cooldown, min confidence, TTL) and
`fusion_math::for_source` (window clamp) replaced with named constants.
- CODE_STYLE inline-path violations (3-segment paths to `BTreeMap`,
`Duration`, `CreateKind`) hoisted to `use`.
Tests: 198 pass (new: 9 fusion_math, 11 attack_type dictionary, 7 adapter
confirmations lookup, 3 model_source serde roundtrip, manifest validator
coverage). \`cargo clippy --package net-guardia -- -D warnings\` clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
231fb88efd
commit
56395f8027
@ -1,7 +1,6 @@
|
||||
# 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.
|
||||
# NetGuardia model manifest. Structural/semantic fields live here;
|
||||
# preprocessing arrays (scaler mean/std, clip params, feature weights) stay
|
||||
# in the JSON sidecar referenced by `preprocessing.scaler_sidecar`.
|
||||
|
||||
name: netguardia-v10
|
||||
adapter: multi_task
|
||||
@ -46,9 +45,12 @@ features:
|
||||
- fwd_bwd_bytes_ratio
|
||||
- iat_cv
|
||||
|
||||
# `confirmations` / `playbook` are parsed today; M2 wires them into the aggregator.
|
||||
# `confirmations` sets the per-class aggregator firing threshold. Classes
|
||||
# with single-shot semantics (C2 / Bot / DNS tunneling / exploit) use 1 so
|
||||
# the aggregator alerts on the first detection; noisier classes can raise
|
||||
# it (DoS/DDoS: 2). Absent entries fall back to the engine default.
|
||||
labels:
|
||||
"0": { name: Bot }
|
||||
"0": { name: Bot, confirmations: 1 }
|
||||
"1": { name: Brute Force }
|
||||
"2": { name: C2 Communication, confirmations: 1 }
|
||||
"3": { name: DNS Tunneling, confirmations: 1 }
|
||||
@ -64,6 +66,9 @@ thresholds:
|
||||
c2: 0.9085615873336792
|
||||
class_min_confidence: 0.4
|
||||
ae: 0.23011694848537445
|
||||
# Average score must exceed `class_min_confidence * alert_multiplier`
|
||||
# before the aggregator fires. Raising this suppresses borderline hits.
|
||||
alert_multiplier: 1.2
|
||||
|
||||
preprocessing:
|
||||
scaler_sidecar: inference_config.json
|
||||
|
||||
135
net-guardia/src/core/detection/fusion_math.rs
Normal file
135
net-guardia/src/core/detection/fusion_math.rs
Normal file
@ -0,0 +1,135 @@
|
||||
//! Fusion policy math primitive — cross-source confidence aggregation.
|
||||
//!
|
||||
//! The policy assumes detection sources are conditionally independent given
|
||||
//! a true attack. In practice ML and CV can share signal on C2 beaconing,
|
||||
//! so a future calibration pass may introduce per-pair weights; this
|
||||
//! module stays the canonical home for whichever formula is in force.
|
||||
|
||||
use crate::model::event::DetectionSource;
|
||||
|
||||
/// Compute `1 − ∏(1 − c_i)` over the given per-source confidences.
|
||||
///
|
||||
/// - Empty input → `0.0` (no evidence).
|
||||
/// - Single input → returns that confidence unchanged.
|
||||
/// - Values are clamped to `[0.0, 1.0]` to keep the result bounded even if
|
||||
/// an upstream source ships noisy unnormalized scores.
|
||||
pub fn fused_confidence(per_source: &[f32]) -> f32 {
|
||||
if per_source.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut inverse: f64 = 1.0;
|
||||
for &c in per_source {
|
||||
let clamped = (c as f64).clamp(0.0, 1.0);
|
||||
inverse *= 1.0 - clamped;
|
||||
}
|
||||
(1.0 - inverse).clamp(0.0, 1.0) as f32
|
||||
}
|
||||
|
||||
/// Default per-source fusion-window length in seconds. Each value scales
|
||||
/// the orchestrator's lookahead budget when that source opens a dedup key.
|
||||
/// Slower sources (Suricata signatures) get longer windows so a follow-up
|
||||
/// ML hit still lands inside; faster sources (ML ticks) use short windows
|
||||
/// because they'd otherwise waste latency waiting on downstream signals.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct FusionWindowLengths {
|
||||
pub suricata_secs: u64,
|
||||
pub cv_secs: u64,
|
||||
pub ml_secs: u64,
|
||||
pub graph_secs: u64,
|
||||
}
|
||||
|
||||
/// Valid range, in seconds, for a fusion window. Clamps protect against a
|
||||
/// misconfigured source opening a wedged (too-long) or useless (zero) key.
|
||||
pub const FUSION_WINDOW_MIN_SECS: u64 = 1;
|
||||
pub const FUSION_WINDOW_MAX_SECS: u64 = 30;
|
||||
|
||||
impl Default for FusionWindowLengths {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
suricata_secs: 10,
|
||||
cv_secs: 8,
|
||||
ml_secs: 2,
|
||||
graph_secs: 5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FusionWindowLengths {
|
||||
/// Range-clamped lookup for the window length of the first source to
|
||||
/// open a fusion key.
|
||||
pub fn for_source(&self, source: DetectionSource) -> u64 {
|
||||
let raw = match source {
|
||||
DetectionSource::Suricata => self.suricata_secs,
|
||||
DetectionSource::Beaconing => self.cv_secs,
|
||||
DetectionSource::ML => self.ml_secs,
|
||||
DetectionSource::Correlation => self.graph_secs,
|
||||
};
|
||||
raw.clamp(FUSION_WINDOW_MIN_SECS, FUSION_WINDOW_MAX_SECS)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn empty_fusion_returns_zero() {
|
||||
assert_eq!(fused_confidence(&[]), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_source_passes_through() {
|
||||
assert!((fused_confidence(&[0.8]) - 0.8).abs() < 1e-5);
|
||||
assert_eq!(fused_confidence(&[0.0]), 0.0);
|
||||
assert_eq!(fused_confidence(&[1.0]), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_sources_boost() {
|
||||
// 1 − (1 − 0.7)(1 − 0.9) = 1 − 0.03 = 0.97
|
||||
let got = fused_confidence(&[0.7, 0.9]);
|
||||
assert!((got - 0.97).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn three_sources_monotone_in_count() {
|
||||
let two = fused_confidence(&[0.6, 0.7]);
|
||||
let three = fused_confidence(&[0.6, 0.7, 0.5]);
|
||||
assert!(three >= two);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn four_sources_bounded() {
|
||||
let four = fused_confidence(&[0.9, 0.8, 0.7, 0.6]);
|
||||
assert!(four < 1.0);
|
||||
assert!(four > 0.99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn out_of_range_confidences_are_clamped() {
|
||||
// Negative / above-1 inputs don't break the math.
|
||||
assert_eq!(fused_confidence(&[-0.5, -0.1]), 0.0);
|
||||
assert_eq!(fused_confidence(&[2.0, 3.0]), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn window_lengths_defaults() {
|
||||
let w = FusionWindowLengths::default();
|
||||
assert_eq!(w.for_source(DetectionSource::Suricata), 10);
|
||||
assert_eq!(w.for_source(DetectionSource::Beaconing), 8);
|
||||
assert_eq!(w.for_source(DetectionSource::ML), 2);
|
||||
assert_eq!(w.for_source(DetectionSource::Correlation), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn window_lengths_clamped_to_range() {
|
||||
let w = FusionWindowLengths {
|
||||
suricata_secs: 999,
|
||||
cv_secs: 0,
|
||||
ml_secs: 2,
|
||||
graph_secs: 5,
|
||||
};
|
||||
assert_eq!(w.for_source(DetectionSource::Suricata), 30);
|
||||
assert_eq!(w.for_source(DetectionSource::Beaconing), 1);
|
||||
}
|
||||
}
|
||||
@ -1,2 +1,3 @@
|
||||
pub mod beaconing;
|
||||
pub mod fusion_math;
|
||||
pub mod orchestrator;
|
||||
|
||||
@ -7,43 +7,65 @@ use macros::log;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::interval;
|
||||
|
||||
use super::fusion_math::{FusionWindowLengths, fused_confidence};
|
||||
use crate::infrastructure::communication_manager::CommunicationManager;
|
||||
use crate::infrastructure::geoip::GeoIpService;
|
||||
use crate::model::detection::attack_type::translate;
|
||||
use crate::model::error::system::SystemError;
|
||||
use crate::model::event::{DetectionEvent, DetectionSource, ThreatDetectedEvent};
|
||||
use crate::model::log::detection::DetectionLog;
|
||||
|
||||
/// Dedup window: detections for the same (source_ip, attack_type) within this window
|
||||
/// are suppressed after the first emission.
|
||||
/// Dedup window: detections for the same `(source_ip, canonical_attack_type)`
|
||||
/// within this window are suppressed after initial fusion-window expiry.
|
||||
const DEDUP_WINDOW_SECS: u64 = 30;
|
||||
|
||||
/// How often to sweep expired dedup entries.
|
||||
const CLEANUP_INTERVAL_SECS: u64 = 60;
|
||||
|
||||
/// Repeat offender detection: same IP within this duration counts as repeat.
|
||||
const REPEAT_OFFENDER_WINDOW_SECS: u64 = 2 * 60 * 60; // 2 hours
|
||||
const REPEAT_OFFENDER_WINDOW_SECS: u64 = 2 * 60 * 60;
|
||||
|
||||
/// Maximum dedup entries to prevent unbounded memory growth under sustained attack.
|
||||
const MAX_DEDUP_ENTRIES: usize = 50_000;
|
||||
|
||||
struct DedupEntry {
|
||||
sources: Vec<DetectionSource>,
|
||||
emitted_at: Instant,
|
||||
/// Per-source record within an in-flight dedup entry. Keeps the strongest
|
||||
/// confidence per source so multi-hit from one source doesn't inflate the
|
||||
/// fused policy.
|
||||
#[derive(Debug, Clone)]
|
||||
struct SourceSample {
|
||||
source: DetectionSource,
|
||||
confidence: f32,
|
||||
}
|
||||
|
||||
/// Coordinates detections from multiple sources (ML, future: rules, correlation, threat feeds).
|
||||
/// Deduplicates, enriches with GeoIP/hit count/repeat offender, and emits ThreatDetectedEvent.
|
||||
struct DedupEntry {
|
||||
sources: Vec<SourceSample>,
|
||||
/// When the orchestrator first emitted for this key.
|
||||
first_emitted_at: Instant,
|
||||
/// Most recent emit (reset on each fused re-emit within the fusion window).
|
||||
emitted_at: Instant,
|
||||
/// Fusion window length for this key, fixed by the first source to arrive.
|
||||
/// Later arrivals don't reset it so the lookahead budget stays predictable.
|
||||
fusion_window: Duration,
|
||||
}
|
||||
|
||||
/// Coordinates detections from ML / Suricata / Beaconing / Correlation.
|
||||
/// Incoming events are canonicalized into a shared attack-type dictionary so
|
||||
/// dedup keys collide across sources; events sharing a key inside the fusion
|
||||
/// window accumulate, and additional sources arriving mid-window trigger a
|
||||
/// re-emit with the combined confidence `1 − ∏(1 − c_i)`.
|
||||
pub struct DetectionOrchestrator {
|
||||
rx: mpsc::Receiver<DetectionEvent>,
|
||||
comm: Arc<CommunicationManager>,
|
||||
geoip: Option<Arc<GeoIpService>>,
|
||||
// Enrichment state
|
||||
// SAFETY: NonZero::new on a literal is infallible.
|
||||
src_ip_counts: lru::LruCache<String, u32>,
|
||||
repeat_tracker: lru::LruCache<String, Instant>,
|
||||
// Dedup state — LRU-bounded to prevent unbounded growth under sustained attack
|
||||
// Dedup state — LRU-bounded to prevent unbounded growth under sustained attack.
|
||||
dedup: lru::LruCache<(String, String), DedupEntry>,
|
||||
dedup_window: Duration,
|
||||
/// Per-source fusion window lengths. Future versions may read overrides
|
||||
/// from DB; the defaults live in `FusionWindowLengths::default`.
|
||||
fusion_windows: FusionWindowLengths,
|
||||
}
|
||||
|
||||
impl DetectionOrchestrator {
|
||||
@ -61,10 +83,10 @@ impl DetectionOrchestrator {
|
||||
repeat_tracker: LruCache::new(NonZero::new(5_000).unwrap()),
|
||||
dedup: LruCache::new(NonZero::new(MAX_DEDUP_ENTRIES).unwrap()),
|
||||
dedup_window: Duration::from_secs(DEDUP_WINDOW_SECS),
|
||||
fusion_windows: FusionWindowLengths::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn the orchestrator as a background task.
|
||||
pub fn start(self) {
|
||||
tokio::spawn(async move { self.run().await });
|
||||
}
|
||||
@ -79,7 +101,7 @@ impl DetectionOrchestrator {
|
||||
event = self.rx.recv() => {
|
||||
match event {
|
||||
Some(detection) => self.handle_detection(detection).await,
|
||||
None => break, // All senders dropped
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = cleanup_interval.tick() => {
|
||||
@ -89,47 +111,115 @@ impl DetectionOrchestrator {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_detection(&mut self, event: DetectionEvent) {
|
||||
async fn handle_detection(&mut self, mut event: DetectionEvent) {
|
||||
// Canonicalize the raw attack_type so Suricata's "brute-force"
|
||||
// classtype and ML's "Brute Force" class name land on the same dedup
|
||||
// key — the precondition for cross-source fusion.
|
||||
let canonical = translate(event.source, &event.attack_type);
|
||||
event.attack_type = canonical.as_str().to_string();
|
||||
|
||||
let key = (event.source_ip.clone(), event.attack_type.clone());
|
||||
let now = Instant::now();
|
||||
|
||||
// Dedup check
|
||||
if let Some(entry) = self.dedup.get(&key)
|
||||
&& now.checked_duration_since(entry.emitted_at).unwrap_or(Duration::ZERO) < self.dedup_window
|
||||
{
|
||||
// Within window: add source attribution but don't re-emit
|
||||
if !entry.sources.contains(&event.source) {
|
||||
// Re-get as mutable to update sources
|
||||
if let Some(entry) = self.dedup.get_mut(&key) {
|
||||
entry.sources.push(event.source.clone());
|
||||
// Path A: existing dedup entry. Decide re-emit (fusion window still
|
||||
// open) vs silence (window closed but dedup still active).
|
||||
if let Some(entry) = self.dedup.get_mut(&key) {
|
||||
let since_first = now.saturating_duration_since(entry.first_emitted_at);
|
||||
|
||||
// Post-dedup-window — treat as a brand-new event (fall through).
|
||||
if since_first >= self.dedup_window {
|
||||
// Expired dedup; fall through to Path B by dropping the entry.
|
||||
self.dedup.pop(&key);
|
||||
} else if since_first < entry.fusion_window {
|
||||
// Still inside the fusion window — accumulate.
|
||||
let is_new_source = !entry.sources.iter().any(|s| s.source == event.source);
|
||||
if is_new_source {
|
||||
entry.sources.push(SourceSample {
|
||||
source: event.source,
|
||||
confidence: event.confidence,
|
||||
});
|
||||
entry.emitted_at = now;
|
||||
} else {
|
||||
// Same source firing again inside the window — keep the
|
||||
// strongest confidence for fusion math.
|
||||
if let Some(existing) = entry.sources.iter_mut().find(|s| s.source == event.source)
|
||||
&& existing.confidence < event.confidence
|
||||
{
|
||||
existing.confidence = event.confidence;
|
||||
}
|
||||
}
|
||||
|
||||
// Only RE-EMIT when a new source joined — same-source
|
||||
// refires are silenced to avoid SOAR cooldown churn.
|
||||
if is_new_source {
|
||||
self.emit_fused(&event, &key).await;
|
||||
} else {
|
||||
log!(DetectionLog::DetectionDeduplicated(
|
||||
event.source_ip.clone(),
|
||||
event.attack_type.clone(),
|
||||
));
|
||||
}
|
||||
return;
|
||||
} else {
|
||||
// Past fusion window, still inside dedup silence → drop.
|
||||
log!(DetectionLog::DetectionDeduplicated(
|
||||
event.source_ip.clone(),
|
||||
event.attack_type.clone(),
|
||||
));
|
||||
return;
|
||||
}
|
||||
log!(DetectionLog::DetectionDeduplicated(event.source_ip, event.attack_type,));
|
||||
return;
|
||||
}
|
||||
|
||||
// Enrich and emit
|
||||
let threat_event = self.enrich(&event).await;
|
||||
let sources = vec![event.source.clone()];
|
||||
|
||||
log!(DetectionLog::DetectionEmitted(
|
||||
event.source_ip.clone(),
|
||||
event.attack_type.clone(),
|
||||
event.confidence,
|
||||
event.ae_score,
|
||||
event.anomaly_score,
|
||||
event.c2_score,
|
||||
sources.len(),
|
||||
));
|
||||
|
||||
// Record dedup entry (LRU-bounded)
|
||||
// Path B: brand-new key (or expired dedup). Emit single-source,
|
||||
// open a fusion window sized by this source.
|
||||
let fusion_window = Duration::from_secs(self.fusion_windows.for_source(event.source));
|
||||
self.dedup.put(
|
||||
key,
|
||||
key.clone(),
|
||||
DedupEntry {
|
||||
sources,
|
||||
sources: vec![SourceSample {
|
||||
source: event.source,
|
||||
confidence: event.confidence,
|
||||
}],
|
||||
first_emitted_at: now,
|
||||
emitted_at: now,
|
||||
fusion_window,
|
||||
},
|
||||
);
|
||||
self.emit_fused(&event, &key).await;
|
||||
}
|
||||
|
||||
/// Build the fused ThreatDetectedEvent from the current dedup entry's
|
||||
/// per-source samples, apply enrichment (hit count / repeat / geoip),
|
||||
/// and publish. Called both on first emit (single source) and on
|
||||
/// within-window re-emit (2..=4 sources).
|
||||
async fn emit_fused(&mut self, trigger_event: &DetectionEvent, key: &(String, String)) {
|
||||
let (sources_vec, fused, per_source_confs): (Vec<DetectionSource>, f32, Vec<f32>) = {
|
||||
let entry = match self.dedup.get(key) {
|
||||
Some(e) => e,
|
||||
None => return,
|
||||
};
|
||||
let confs: Vec<f32> = entry.sources.iter().map(|s| s.confidence).collect();
|
||||
let sources: Vec<DetectionSource> = entry.sources.iter().map(|s| s.source).collect();
|
||||
let fused = fused_confidence(&confs);
|
||||
(sources, fused, confs)
|
||||
};
|
||||
let _ = per_source_confs; // currently only used for the log line below
|
||||
|
||||
let mut threat_event = self.enrich(trigger_event).await;
|
||||
threat_event.sources = sources_vec;
|
||||
threat_event.active_source_count = threat_event.sources.len();
|
||||
threat_event.fused_confidence = fused;
|
||||
threat_event.confidence = fused;
|
||||
|
||||
log!(DetectionLog::DetectionEmitted(
|
||||
trigger_event.source_ip.clone(),
|
||||
trigger_event.attack_type.clone(),
|
||||
threat_event.confidence,
|
||||
trigger_event.ae_score,
|
||||
trigger_event.anomaly_score,
|
||||
trigger_event.c2_score,
|
||||
threat_event.active_source_count,
|
||||
));
|
||||
|
||||
if let Err(e) = self.comm.publish_event(threat_event).await {
|
||||
log!(SystemError::MlSoarBridgeFailed(e));
|
||||
@ -139,14 +229,12 @@ impl DetectionOrchestrator {
|
||||
async fn enrich(&mut self, event: &DetectionEvent) -> ThreatDetectedEvent {
|
||||
let src_ip = &event.source_ip;
|
||||
|
||||
// Compute packet rate
|
||||
let packet_rate = if event.flow_duration_us > 0 {
|
||||
event.packet_count as f64 / (event.flow_duration_us as f64 / 1_000_000.0)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Update hit count (LRU bounded)
|
||||
let hit_count = match self.src_ip_counts.get_mut(src_ip) {
|
||||
Some(c) => {
|
||||
*c = c.saturating_add(1);
|
||||
@ -158,7 +246,6 @@ impl DetectionOrchestrator {
|
||||
}
|
||||
};
|
||||
|
||||
// Check repeat offender (same IP within window)
|
||||
let repeat_window = Duration::from_secs(REPEAT_OFFENDER_WINDOW_SECS);
|
||||
let now = Instant::now();
|
||||
let is_repeat = self
|
||||
@ -167,7 +254,6 @@ impl DetectionOrchestrator {
|
||||
.is_some_and(|last| now.checked_duration_since(*last).unwrap_or(Duration::ZERO) < repeat_window);
|
||||
self.repeat_tracker.put(src_ip.clone(), now);
|
||||
|
||||
// GeoIP lookup
|
||||
let geoip_country = if let Some(ref svc) = self.geoip {
|
||||
if let Ok(ip) = src_ip.parse() {
|
||||
svc.lookup(ip).await.ok().flatten().and_then(|loc| loc.country_code)
|
||||
@ -188,7 +274,12 @@ impl DetectionOrchestrator {
|
||||
protocol: event.protocol,
|
||||
geoip_country,
|
||||
is_repeat_offender: is_repeat,
|
||||
sources: vec![event.source.clone()],
|
||||
// These three get overwritten in `emit_fused` with the
|
||||
// accumulated values; initialize to the single-source defaults
|
||||
// so a direct caller also gets a consistent shape.
|
||||
sources: vec![event.source],
|
||||
active_source_count: 1,
|
||||
fused_confidence: event.confidence,
|
||||
ae_score: event.ae_score,
|
||||
anomaly_score: event.anomaly_score,
|
||||
c2_score: event.c2_score,
|
||||
@ -198,9 +289,12 @@ impl DetectionOrchestrator {
|
||||
fn cleanup_expired(&mut self) {
|
||||
let now = Instant::now();
|
||||
let window = self.dedup_window;
|
||||
// Pop expired entries from the LRU (oldest entries are least recently used)
|
||||
while let Some((_, entry)) = self.dedup.peek_lru() {
|
||||
if now.checked_duration_since(entry.emitted_at).unwrap_or(Duration::ZERO) >= window {
|
||||
if now
|
||||
.checked_duration_since(entry.first_emitted_at)
|
||||
.unwrap_or(Duration::ZERO)
|
||||
>= window
|
||||
{
|
||||
self.dedup.pop_lru();
|
||||
} else {
|
||||
break;
|
||||
@ -208,3 +302,11 @@ impl DetectionOrchestrator {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
//! Orchestrator integration tests go here once we have an in-memory
|
||||
//! CommunicationManager harness. The fusion math is covered by
|
||||
//! `fusion_math::tests` and canonical translation by
|
||||
//! `model::detection::attack_type::tests`.
|
||||
}
|
||||
|
||||
235
net-guardia/src/core/ml/adapter.rs
Normal file
235
net-guardia/src/core/ml/adapter.rs
Normal file
@ -0,0 +1,235 @@
|
||||
//! `MLModelAdapter` — the three ONNX adapter shapes the inference pipeline
|
||||
//! understands, plus the wrapping `ModelSourceState` machine used inside
|
||||
//! `Inference::models: ArcSwap<ModelSourceState>`.
|
||||
//!
|
||||
//! Day 1 is `Dormant` (no model loaded); the happy path is `Active { adapter,
|
||||
//! info }`; failed loads park in `Error { msg, since, last_attempted_path }`.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use super::manifest::LabelSpec;
|
||||
use crate::model::detection::ml_detection::RunnableModel;
|
||||
use crate::model::detection::model_source::{ModelInfo, ModelSourceStatus};
|
||||
|
||||
/// Compile-time sanity: `RunnableModel` must be `Send + Sync` because we
|
||||
/// stuff it inside an `ArcSwap`. If a future `tract-onnx` upgrade silently
|
||||
/// drops the bounds, this line stops compiling and we catch it before it
|
||||
/// becomes a production data race.
|
||||
const _: fn() = || {
|
||||
fn assert_send_sync<T: Send + Sync>() {}
|
||||
assert_send_sync::<RunnableModel>();
|
||||
};
|
||||
|
||||
/// One of three adapter shapes. Each variant carries the RunnableModel(s)
|
||||
/// it needs plus the metadata required to dispatch inference:
|
||||
/// batch size, feature counts, label map (classifier / multi-task only).
|
||||
#[derive(Clone)]
|
||||
pub enum MLModelAdapter {
|
||||
/// Pure anomaly-detection autoencoder. Output is a single MSE-style
|
||||
/// reconstruction error per flow; classifier-less. v1 product main path.
|
||||
AutoencoderOnly {
|
||||
model: Arc<RunnableModel>,
|
||||
batch_size: usize,
|
||||
n_features: usize,
|
||||
},
|
||||
/// Single classifier, no AE. Output is a per-class softmax tensor.
|
||||
/// Labels come from the manifest's `labels` map; `normal_idx` is
|
||||
/// pre-resolved so `infer_batch` doesn't re-scan on every tick.
|
||||
ClassifierOnly {
|
||||
model: Arc<RunnableModel>,
|
||||
batch_size: usize,
|
||||
n_features: usize,
|
||||
labels: BTreeMap<String, LabelSpec>,
|
||||
normal_idx: Option<usize>,
|
||||
},
|
||||
/// AE + classifier in one adapter. Classifier consumes (ae_features, ae_score).
|
||||
/// Three output tensors: anomaly_score / class_probs / c2_score.
|
||||
MultiTask {
|
||||
ae: Arc<RunnableModel>,
|
||||
classifier: Arc<RunnableModel>,
|
||||
batch_size: usize,
|
||||
n_ae: usize,
|
||||
n_cls: usize,
|
||||
labels: BTreeMap<String, LabelSpec>,
|
||||
normal_idx: Option<usize>,
|
||||
c2_idx: Option<usize>,
|
||||
},
|
||||
}
|
||||
|
||||
impl MLModelAdapter {
|
||||
/// Per-attack-type confirmations lookup: resolve the manifest's
|
||||
/// `labels[*].confirmations` value whose label name matches
|
||||
/// `attack_type_name` (case-insensitive). Callers fall back to their own
|
||||
/// default when this returns `None` — `AutoencoderOnly` carries no labels,
|
||||
/// and classifier labels may legitimately omit the override.
|
||||
pub fn confirmations_for(&self, attack_type_name: &str) -> Option<usize> {
|
||||
match self {
|
||||
Self::MultiTask { labels, .. } | Self::ClassifierOnly { labels, .. } => {
|
||||
confirmations_from_labels(labels, attack_type_name)
|
||||
}
|
||||
Self::AutoencoderOnly { .. } => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pure-data lookup extracted so unit tests can cover the matching rules
|
||||
/// without constructing a full `MLModelAdapter` (which requires a real
|
||||
/// `RunnableModel` — expensive and fragile to mock).
|
||||
fn confirmations_from_labels(labels: &BTreeMap<String, LabelSpec>, attack_type_name: &str) -> Option<usize> {
|
||||
labels
|
||||
.values()
|
||||
.find(|spec| spec.name.eq_ignore_ascii_case(attack_type_name))
|
||||
.and_then(|spec| spec.confirmations)
|
||||
}
|
||||
|
||||
/// Internal state machine. Held inside `ArcSwap<ModelSourceState>` so the
|
||||
/// inference pipeline can check state once per tick without locks and
|
||||
/// atomically swap on upload / deletion / reload failure.
|
||||
///
|
||||
/// Variants are deliberately **not** `Serialize` — the `Active` variant
|
||||
/// holds an `Arc<RunnableModel>` which isn't serde-friendly. Callers that
|
||||
/// need a wire representation call `to_status()` to produce the lightweight
|
||||
/// `ModelSourceStatus` consumed by WebSocket / HTTP.
|
||||
pub enum ModelSourceState {
|
||||
/// Day 1 default. `models/` empty, or admin deleted the current model.
|
||||
/// Drift detector becomes a no-op; fusion math still runs with the
|
||||
/// remaining 3 sources.
|
||||
Dormant,
|
||||
/// Model loaded, inference ticks consume flows.
|
||||
Active { adapter: MLModelAdapter, info: ModelInfo },
|
||||
/// Last load attempt failed — schema mismatch, timeout, corrupted
|
||||
/// ONNX. Inference skips; UI renders the reason. Swap to Dormant /
|
||||
/// Active via normal reload path.
|
||||
Error {
|
||||
msg: String,
|
||||
since: SystemTime,
|
||||
last_attempted_path: Option<PathBuf>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ModelSourceState {
|
||||
pub fn is_active(&self) -> bool {
|
||||
matches!(self, Self::Active { .. })
|
||||
}
|
||||
|
||||
/// Wire-format snapshot for UI / HTTP. Never borrows the adapter — the
|
||||
/// returned value is safe to send across WebSocket boundaries.
|
||||
pub fn to_status(&self) -> ModelSourceStatus {
|
||||
match self {
|
||||
Self::Dormant => ModelSourceStatus::Dormant,
|
||||
Self::Active { info, .. } => ModelSourceStatus::Active { info: info.clone() },
|
||||
Self::Error {
|
||||
msg,
|
||||
since,
|
||||
last_attempted_path,
|
||||
} => ModelSourceStatus::Error {
|
||||
msg: msg.clone(),
|
||||
since_secs: since
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0),
|
||||
last_attempted_path: last_attempted_path.as_ref().map(|p| p.display().to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn dormant_state_serializes_to_dormant_status() {
|
||||
let state = ModelSourceState::Dormant;
|
||||
let status = state.to_status();
|
||||
assert!(status.is_dormant());
|
||||
}
|
||||
|
||||
fn sample_labels() -> BTreeMap<String, LabelSpec> {
|
||||
let mut m = BTreeMap::new();
|
||||
m.insert(
|
||||
"0".into(),
|
||||
LabelSpec {
|
||||
name: "Bot".into(),
|
||||
confirmations: Some(1),
|
||||
playbook: None,
|
||||
},
|
||||
);
|
||||
m.insert(
|
||||
"4".into(),
|
||||
LabelSpec {
|
||||
name: "DoS/DDoS".into(),
|
||||
confirmations: Some(2),
|
||||
playbook: None,
|
||||
},
|
||||
);
|
||||
m.insert(
|
||||
"7".into(),
|
||||
LabelSpec {
|
||||
name: "Normal".into(),
|
||||
confirmations: None,
|
||||
playbook: None,
|
||||
},
|
||||
);
|
||||
m
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmations_lookup_hits_exact_name() {
|
||||
let labels = sample_labels();
|
||||
assert_eq!(confirmations_from_labels(&labels, "Bot"), Some(1));
|
||||
assert_eq!(confirmations_from_labels(&labels, "DoS/DDoS"), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmations_lookup_is_case_insensitive() {
|
||||
let labels = sample_labels();
|
||||
assert_eq!(confirmations_from_labels(&labels, "bot"), Some(1));
|
||||
assert_eq!(confirmations_from_labels(&labels, "dos/ddos"), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmations_lookup_label_without_override_returns_none() {
|
||||
let labels = sample_labels();
|
||||
assert_eq!(confirmations_from_labels(&labels, "Normal"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmations_lookup_unknown_name_returns_none() {
|
||||
let labels = sample_labels();
|
||||
assert_eq!(confirmations_from_labels(&labels, "Phantom"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmations_lookup_empty_labels_returns_none() {
|
||||
let labels = BTreeMap::new();
|
||||
assert_eq!(confirmations_from_labels(&labels, "Bot"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_state_preserves_details() {
|
||||
let state = ModelSourceState::Error {
|
||||
msg: "shape mismatch".to_string(),
|
||||
since: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
|
||||
last_attempted_path: Some(PathBuf::from("models/bad.onnx")),
|
||||
};
|
||||
let status = state.to_status();
|
||||
match status {
|
||||
ModelSourceStatus::Error {
|
||||
msg,
|
||||
since_secs,
|
||||
last_attempted_path,
|
||||
} => {
|
||||
assert_eq!(msg, "shape mismatch");
|
||||
assert_eq!(since_secs, 1_700_000_000);
|
||||
assert_eq!(last_attempted_path.as_deref(), Some("models/bad.onnx"));
|
||||
}
|
||||
_ => panic!("expected Error"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,3 +1,9 @@
|
||||
//! Per-flow detection aggregator.
|
||||
//!
|
||||
//! Records per-key hits within a rolling time window; callers decide when to
|
||||
//! fire based on how many hits a given attack class needs (manifest-driven)
|
||||
//! and how far the rolling-average score beats the confidence threshold.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@ -6,40 +12,40 @@ use crate::model::detection::ml_detection::FlowKey;
|
||||
pub struct AttackAggregator {
|
||||
detections: HashMap<FlowKey, Vec<(Instant, f32)>>,
|
||||
window_duration: Duration,
|
||||
min_detections: usize,
|
||||
alert_threshold_multiplier: f32,
|
||||
}
|
||||
|
||||
impl AttackAggregator {
|
||||
pub fn new(window_secs: u64, min_detections: usize) -> Self {
|
||||
pub fn new(window_secs: u64) -> Self {
|
||||
Self {
|
||||
detections: HashMap::new(),
|
||||
window_duration: Duration::from_secs(window_secs),
|
||||
min_detections,
|
||||
alert_threshold_multiplier: 1.2,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn should_alert(&mut self, flow_key: &FlowKey, score: f32, threshold: f32, attack_type: Option<&str>) -> bool {
|
||||
/// Record a detection and decide whether the flow should fire an alert.
|
||||
///
|
||||
/// - `required_confirmations`: in-window hit count the flow must reach.
|
||||
/// Resolved by the caller from the active manifest's per-label value;
|
||||
/// validated at manifest load to be ≥ 1, so no runtime floor is needed.
|
||||
/// - `alert_multiplier`: scales `threshold` before the average-score
|
||||
/// comparison, driven by the manifest's `thresholds.alert_multiplier`.
|
||||
pub fn should_alert(
|
||||
&mut self,
|
||||
flow_key: &FlowKey,
|
||||
score: f32,
|
||||
threshold: f32,
|
||||
required_confirmations: usize,
|
||||
alert_multiplier: f32,
|
||||
) -> bool {
|
||||
let now = Instant::now();
|
||||
|
||||
let detections = self.detections.entry(flow_key.clone()).or_default();
|
||||
detections.retain(|(time, _)| now.duration_since(*time) < self.window_duration);
|
||||
detections.push((now, score));
|
||||
|
||||
// Per-attack-type adaptive min_detections:
|
||||
// DoS/DDoS: high frequency, need more confirmations to avoid alert storms
|
||||
// C2/Bot/DNS Tunneling/Exploitation: low frequency or single-shot, alert immediately
|
||||
let effective_min = match attack_type {
|
||||
Some("DoS/DDoS") => self.min_detections.saturating_mul(2).max(1),
|
||||
Some("C2 Communication") | Some("Bot") | Some("DNS Tunneling") | Some("Exploitation") => 1,
|
||||
_ => self.min_detections,
|
||||
};
|
||||
|
||||
if detections.len() >= effective_min {
|
||||
if detections.len() >= required_confirmations {
|
||||
let avg_score: f32 = detections.iter().map(|(_, s)| s).sum::<f32>() / detections.len() as f32;
|
||||
|
||||
return avg_score > threshold * self.alert_threshold_multiplier;
|
||||
return avg_score > threshold * alert_multiplier;
|
||||
}
|
||||
|
||||
false
|
||||
@ -58,6 +64,8 @@ impl AttackAggregator {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const TEST_MULTIPLIER: f32 = 1.2;
|
||||
|
||||
fn test_key() -> FlowKey {
|
||||
FlowKey {
|
||||
src_ip: [192, 168, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
@ -70,61 +78,66 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_attack_type_uses_base_min_detections() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
fn required_three_fires_on_third_hit() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
// Need 3 detections for default type
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dos_ddos_requires_double_min_detections() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
fn required_one_fires_immediately() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, 1, TEST_MULTIPLIER));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_six_needs_six_hits() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
// DoS/DDoS needs 6 detections (3 * 2)
|
||||
for _ in 0..5 {
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("DoS/DDoS")));
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, 6, TEST_MULTIPLIER));
|
||||
}
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("DoS/DDoS")));
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, 6, TEST_MULTIPLIER));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn c2_alerts_on_first_detection() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
fn average_score_at_or_below_scaled_threshold_does_not_fire() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
// C2 Communication alerts immediately (min=1)
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("C2 Communication")));
|
||||
// score 1.0, threshold 1.0, multiplier 1.2 → gate is 1.2; 1.0 misses.
|
||||
assert!(!agg.should_alert(&key, 1.0, 1.0, 1, TEST_MULTIPLIER));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exploitation_alerts_on_first_detection() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
fn larger_multiplier_raises_the_bar() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Exploitation")));
|
||||
// multiplier 2.5, threshold 1.0 → gate is 2.5; score 2.0 misses.
|
||||
assert!(!agg.should_alert(&key, 2.0, 1.0, 1, 2.5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bot_alerts_on_first_detection() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
fn cleanup_preserves_fresh_entries() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key = test_key();
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Bot")));
|
||||
agg.should_alert(&key, 5.0, 1.0, 10, TEST_MULTIPLIER);
|
||||
assert!(agg.detections.contains_key(&key));
|
||||
agg.cleanup();
|
||||
assert!(agg.detections.contains_key(&key));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dns_tunneling_alerts_on_first_detection() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
let key = test_key();
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, Some("DNS Tunneling")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_attack_type_uses_default() {
|
||||
let mut agg = AttackAggregator::new(60, 3);
|
||||
let key = test_key();
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, None));
|
||||
assert!(!agg.should_alert(&key, 5.0, 1.0, None));
|
||||
assert!(agg.should_alert(&key, 5.0, 1.0, None));
|
||||
fn independent_flows_track_separately() {
|
||||
let mut agg = AttackAggregator::new(60);
|
||||
let key_a = test_key();
|
||||
let mut key_b = test_key();
|
||||
key_b.dst_port = 81;
|
||||
assert!(!agg.should_alert(&key_a, 5.0, 1.0, 2, TEST_MULTIPLIER));
|
||||
assert!(!agg.should_alert(&key_b, 5.0, 1.0, 2, TEST_MULTIPLIER));
|
||||
assert!(agg.should_alert(&key_a, 5.0, 1.0, 2, TEST_MULTIPLIER));
|
||||
assert!(agg.should_alert(&key_b, 5.0, 1.0, 2, TEST_MULTIPLIER));
|
||||
}
|
||||
}
|
||||
|
||||
@ -95,12 +95,9 @@ fn reconcile_features(
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
@ -113,6 +110,9 @@ fn apply_manifest_overrides(manifest: &ModelManifest, config: &mut MLInferenceCo
|
||||
if let Some(v) = manifest.thresholds.ae {
|
||||
config.ae_threshold = v;
|
||||
}
|
||||
if let Some(v) = manifest.thresholds.alert_multiplier {
|
||||
config.alert_threshold_multiplier = v;
|
||||
}
|
||||
}
|
||||
|
||||
fn manifest_labels_to_map(labels: &BTreeMap<String, LabelSpec>) -> HashMap<String, String> {
|
||||
@ -166,6 +166,7 @@ mod tests {
|
||||
anomaly_threshold: 0.5,
|
||||
c2_threshold: 0.5,
|
||||
class_min_confidence: 0.4,
|
||||
alert_threshold_multiplier: 1.2,
|
||||
model_type: "MultiTaskModel".into(),
|
||||
output_names: vec!["anomaly".into(), "class_probs".into(), "c2_score".into()],
|
||||
ae_feature_weights: HashMap::new(),
|
||||
|
||||
@ -1,3 +1,9 @@
|
||||
//! ML engine — orchestrates flow tracking, feature extraction, adapter-
|
||||
//! dispatched inference, and alert broadcast. Each inference tick cleans up
|
||||
//! stale flows, gathers the latest batch, optionally logs a Flow Trace row,
|
||||
//! and — when the inference source is Active — updates drift and runs
|
||||
//! adapter-dispatched inference.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
@ -12,14 +18,17 @@ use super::alert::MLAlert;
|
||||
use super::drift_detector::DriftDetector;
|
||||
use super::flow_tracker::{FlowData, FlowTracker};
|
||||
use super::inference::Inference;
|
||||
use super::model_loader::MLModels;
|
||||
use super::traffic_logger::TrafficLogger;
|
||||
use crate::interface::port::packet_sink::{PacketSink, PacketSinkFactory};
|
||||
use crate::model::detection::flow_features::FlowFeatures;
|
||||
use crate::model::detection::ml_detection::{EngineConfig, FlowKey, InferenceStats};
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::monitoring::user_packet::UserPacket;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
|
||||
/// Divisor applied to "ticks per aggregator window" to derive the fallback
|
||||
/// confirmations count: 2 means a default-behaviour detection must fire
|
||||
/// across at least half the window's ticks before alerting.
|
||||
const DEFAULT_CONFIRMATION_WINDOW_FRACTION: u64 = 2;
|
||||
|
||||
/// Per-queue tracker. With symmetric hash in eBPF, both directions of a flow
|
||||
/// land on the same queue, so per-queue trackers correctly see bidirectional flows.
|
||||
@ -32,29 +41,29 @@ pub struct Engine {
|
||||
drift_detector: Arc<Mutex<DriftDetector>>,
|
||||
ml_alert: Arc<MLAlert>,
|
||||
min_packets: usize,
|
||||
/// Confirmations count used when the active manifest's label has no
|
||||
/// explicit `confirmations` override.
|
||||
default_confirmations: usize,
|
||||
batch_size: usize,
|
||||
inference_interval_secs: u64,
|
||||
traffic_logger: Option<Arc<TrafficLogger>>,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
/// Build an Engine around an already-constructed Inference pipeline.
|
||||
/// The Inference's state (Dormant / Active / Error) is consulted per tick.
|
||||
pub fn new(
|
||||
models: Arc<MLModels>,
|
||||
config: Arc<MLInferenceConfig>,
|
||||
inference_pipeline: Arc<Inference>,
|
||||
ml_alert: Arc<MLAlert>,
|
||||
drift_detector: Arc<Mutex<DriftDetector>>,
|
||||
engine_config: EngineConfig,
|
||||
traffic_logger: Option<Arc<TrafficLogger>>,
|
||||
num_threads: u32,
|
||||
) -> Self {
|
||||
let inference_pipeline = Arc::new(Inference::new(models, config));
|
||||
|
||||
let min_detections =
|
||||
((engine_config.aggregator_window_secs / engine_config.inference_interval_secs) / 2).max(1) as usize;
|
||||
let aggregator = Mutex::new(AttackAggregator::new(
|
||||
engine_config.aggregator_window_secs,
|
||||
min_detections,
|
||||
));
|
||||
let interval_secs = engine_config.inference_interval_secs.max(1);
|
||||
let ticks_per_window = engine_config.aggregator_window_secs / interval_secs;
|
||||
let default_confirmations = (ticks_per_window / DEFAULT_CONFIRMATION_WINDOW_FRACTION).max(1) as usize;
|
||||
let aggregator = Mutex::new(AttackAggregator::new(engine_config.aggregator_window_secs));
|
||||
|
||||
let max_flows_per_thread = engine_config.max_flows / (num_threads as usize).max(1);
|
||||
let trackers: Vec<ThreadTracker> = (0..num_threads)
|
||||
@ -68,16 +77,13 @@ impl Engine {
|
||||
drift_detector,
|
||||
ml_alert,
|
||||
min_packets: engine_config.min_packets,
|
||||
default_confirmations,
|
||||
batch_size: engine_config.batch_size,
|
||||
inference_interval_secs: engine_config.inference_interval_secs,
|
||||
traffic_logger,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn inference_pipeline(&self) -> &Arc<Inference> {
|
||||
&self.inference_pipeline
|
||||
}
|
||||
|
||||
/// xsk_manager calls this per queue_id; with symmetric hash each queue has its own tracker.
|
||||
pub fn tracker(&self, queue_id: u32) -> &ThreadTracker {
|
||||
&self.trackers[queue_id as usize % self.trackers.len()]
|
||||
@ -103,21 +109,15 @@ impl Engine {
|
||||
1 => 1,
|
||||
// UDP
|
||||
17 => match flow_key.dst_port {
|
||||
// DNS: amplification, DGA, tunneling can be 1-2 packets
|
||||
53 => 1,
|
||||
// NTP amplification
|
||||
123 => 2,
|
||||
// Known mining pool ports
|
||||
3333 | 45700 => 2,
|
||||
_ => global,
|
||||
},
|
||||
// TCP
|
||||
6 => match flow_key.dst_port {
|
||||
// DNS over TCP
|
||||
53 => 2,
|
||||
// Common C2 ports: Metasploit, Cobalt Strike, reverse shells
|
||||
4444 | 8443 | 8080 | 1337 | 31337 => 2,
|
||||
// Stratum mining
|
||||
3333 | 45700 => 2,
|
||||
_ => global,
|
||||
},
|
||||
@ -142,7 +142,6 @@ impl Engine {
|
||||
_ = ticker.tick() => {}
|
||||
}
|
||||
|
||||
// Move CPU-bound ML inference off the tokio executor
|
||||
let engine = Arc::clone(&self);
|
||||
let _ = spawn_blocking(move || {
|
||||
engine.run_inference_tick();
|
||||
@ -151,6 +150,10 @@ impl Engine {
|
||||
}
|
||||
}
|
||||
|
||||
/// Phased tick: clean up stale flows, gather the uninferred batch
|
||||
/// (min-packet filtered), optionally write a Flow Trace row, and when
|
||||
/// the inference source is Active update drift then run inference.
|
||||
/// Dormant / Error states short-circuit before drift + inference.
|
||||
fn run_inference_tick(&self) {
|
||||
let mut all_flows = Vec::new();
|
||||
let mut total_count = 0;
|
||||
@ -159,13 +162,11 @@ impl Engine {
|
||||
.map(|d| d.as_micros() as u64)
|
||||
.unwrap_or(0);
|
||||
|
||||
// Phase 0: clean up stale / terminated flows
|
||||
for tracker in &self.trackers {
|
||||
let mut t = tracker.lock();
|
||||
t.cleanup_stale_flows(now_us);
|
||||
}
|
||||
|
||||
// Phase 1: short lock per tracker — clone uninferred flows, mark as inferred
|
||||
for tracker in &self.trackers {
|
||||
let mut t = tracker.lock();
|
||||
total_count += t.flow_count();
|
||||
@ -174,7 +175,6 @@ impl Engine {
|
||||
flow.packet_count() >= Self::effective_min_packets(&flow.flow_key, self.min_packets)
|
||||
}),
|
||||
);
|
||||
// lock released here
|
||||
}
|
||||
|
||||
log!(MLLog::FlowStats(
|
||||
@ -190,9 +190,14 @@ impl Engine {
|
||||
|
||||
if let Some(ref logger) = self.traffic_logger {
|
||||
self.log_traffic(&all_flows, logger);
|
||||
} else {
|
||||
self.run_inference(&all_flows);
|
||||
}
|
||||
|
||||
if !self.inference_pipeline.is_active() {
|
||||
return;
|
||||
}
|
||||
|
||||
self.update_drift(&all_flows);
|
||||
self.run_inference(&all_flows);
|
||||
}
|
||||
|
||||
fn log_traffic(&self, flows: &[FlowData], logger: &TrafficLogger) {
|
||||
@ -203,32 +208,36 @@ impl Engine {
|
||||
}
|
||||
}
|
||||
|
||||
/// Update drift baseline — only called when inference state is Active,
|
||||
/// guaranteeing the ArcSwap snapshot the next `infer_batch` sees matches
|
||||
/// the features we just normalized against.
|
||||
fn update_drift(&self, batch: &[FlowData]) {
|
||||
let batch = &batch[..batch.len().min(self.batch_size)];
|
||||
let config = &self.inference_pipeline.config;
|
||||
let mut dd = self.drift_detector.lock();
|
||||
for flow in batch {
|
||||
let features = FlowFeatures::extract(flow, &config.ae_feature_names);
|
||||
let normalized: Vec<f64> = features
|
||||
.features
|
||||
.iter()
|
||||
.zip(config.ae_scaler_mean.iter().zip(config.ae_scaler_std.iter()))
|
||||
.map(|(&val, (&mean, &std))| if std.abs() > 1e-12 { (val - mean) / std } else { 0.0 })
|
||||
.collect();
|
||||
dd.update(&normalized);
|
||||
}
|
||||
}
|
||||
|
||||
fn run_inference(&self, flows: &[FlowData]) {
|
||||
let batch = &flows[..flows.len().min(self.batch_size)];
|
||||
|
||||
log!(MLLog::RunningInference(batch.len()));
|
||||
|
||||
// Feed normalized features into drift detector for each flow in the batch
|
||||
{
|
||||
let config = &self.inference_pipeline.config;
|
||||
let mut dd = self.drift_detector.lock();
|
||||
for flow in batch.iter() {
|
||||
let features = FlowFeatures::extract(flow, &config.ae_feature_names);
|
||||
let normalized: Vec<f64> = features
|
||||
.features
|
||||
.iter()
|
||||
.zip(config.ae_scaler_mean.iter().zip(config.ae_scaler_std.iter()))
|
||||
.map(|(&val, (&mean, &std))| if std.abs() > 1e-12 { (val - mean) / std } else { 0.0 })
|
||||
.collect();
|
||||
dd.update(&normalized);
|
||||
}
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let results = self.inference_pipeline.infer_batch(batch);
|
||||
let elapsed_us = start.elapsed().as_micros() as u64;
|
||||
|
||||
let stats = InferenceStats::from_results(&results, elapsed_us);
|
||||
self.inference_pipeline.record_tick_qps(stats.flows_per_second);
|
||||
|
||||
if results.len() != batch.len() {
|
||||
log!(MLLog::InferenceResults(batch.len(), results.len()));
|
||||
@ -242,39 +251,42 @@ impl Engine {
|
||||
stats.flows_per_second
|
||||
));
|
||||
|
||||
{
|
||||
let mut aggregator = self.aggregator.lock();
|
||||
for result in &results {
|
||||
if result.is_attack {
|
||||
let should_alert = aggregator.should_alert(
|
||||
&result.flow_key_raw,
|
||||
let mut aggregator = self.aggregator.lock();
|
||||
let config = &self.inference_pipeline.config;
|
||||
for result in &results {
|
||||
if result.is_attack {
|
||||
let required_confirmations = result
|
||||
.attack_type
|
||||
.as_deref()
|
||||
.and_then(|at| self.inference_pipeline.confirmations_for_attack_type(at))
|
||||
.unwrap_or(self.default_confirmations);
|
||||
let should_alert = aggregator.should_alert(
|
||||
&result.flow_key_raw,
|
||||
result.confidence,
|
||||
config.class_min_confidence,
|
||||
required_confirmations,
|
||||
config.alert_threshold_multiplier,
|
||||
);
|
||||
|
||||
if should_alert {
|
||||
log!(MLLog::ThreatDetected(
|
||||
format!("{:?}", result.direction),
|
||||
result.flow_key.clone(),
|
||||
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
|
||||
result.confidence,
|
||||
self.inference_pipeline.config.class_min_confidence,
|
||||
result.attack_type.as_deref(),
|
||||
);
|
||||
result.ae_score,
|
||||
));
|
||||
|
||||
if should_alert {
|
||||
log!(MLLog::ThreatDetected(
|
||||
format!("{:?}", result.direction),
|
||||
result.flow_key.clone(),
|
||||
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
|
||||
result.confidence,
|
||||
result.ae_score,
|
||||
));
|
||||
|
||||
self.ml_alert.broadcast_alert(result);
|
||||
}
|
||||
self.ml_alert.broadcast_alert(result);
|
||||
}
|
||||
}
|
||||
|
||||
aggregator.cleanup();
|
||||
}
|
||||
|
||||
aggregator.cleanup();
|
||||
}
|
||||
}
|
||||
|
||||
/// Adapter that exposes one `ThreadTracker` (per AF_XDP queue) as a
|
||||
/// `PacketSink`. `XskManager` holds `Arc<dyn PacketSink>` per queue and never
|
||||
/// touches `FlowTracker` concrete types.
|
||||
/// Adapter that exposes one `ThreadTracker` (per AF_XDP queue) as a `PacketSink`.
|
||||
struct QueueTrackerSink {
|
||||
tracker: ThreadTracker,
|
||||
}
|
||||
|
||||
@ -1,4 +1,16 @@
|
||||
//! Inference pipeline — state-aware, adapter-dispatched.
|
||||
//!
|
||||
//! The pipeline holds `ArcSwap<ModelSourceState>` so the engine can take a
|
||||
//! lock-free snapshot per tick and `infer_batch` dispatches on whichever
|
||||
//! `MLModelAdapter` variant the active state carries.
|
||||
//!
|
||||
//! The MultiTask arm fires on any of: anomaly head above threshold,
|
||||
//! classifier picking a non-Normal class at ≥ `class_min_confidence`, or the
|
||||
//! C2 head elevated above its own threshold — so AE reconstruction error
|
||||
//! can't mask a stealthy attack the classifier does recognize.
|
||||
|
||||
use std::cmp::Ordering as CmpOrdering;
|
||||
use std::collections::BTreeMap;
|
||||
use std::panic::{self, AssertUnwindSafe};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
@ -8,14 +20,16 @@ use arc_swap::ArcSwap;
|
||||
use macros::log;
|
||||
use tract_onnx::prelude::*;
|
||||
|
||||
use super::adapter::{MLModelAdapter, ModelSourceState};
|
||||
use super::flow_tracker::FlowData;
|
||||
use super::model_loader::MLModels;
|
||||
use super::manifest::LabelSpec;
|
||||
use crate::model::detection::flow_features::FlowFeatures;
|
||||
use crate::model::detection::ml_detection::DetectionResult;
|
||||
use crate::model::detection::ml_detection::{DetectionResult, RunnableModel};
|
||||
use crate::model::detection::model_source::ModelSourceStatus;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
|
||||
/// (anomaly_scores, per_class_probs, c2_scores)
|
||||
/// (anomaly_scores, per_class_probs, c2_scores) — MultiTask batch output.
|
||||
type ClassifierBatchOutput = (Vec<f32>, Vec<Vec<f32>>, Vec<f32>);
|
||||
|
||||
/// Consecutive failures to trip the circuit breaker.
|
||||
@ -26,56 +40,87 @@ const CIRCUIT_BREAKER_WINDOW_SECS: u64 = 60;
|
||||
const CIRCUIT_BREAKER_COOLDOWN_SECS: u64 = 120;
|
||||
|
||||
pub struct Inference {
|
||||
models: ArcSwap<MLModels>,
|
||||
state: ArcSwap<ModelSourceState>,
|
||||
pub config: Arc<MLInferenceConfig>,
|
||||
c2_class_idx: Option<usize>,
|
||||
normal_class_idx: Option<usize>,
|
||||
/// Rolling-window QPS estimate, published in `ModelInfo.qps_recent`.
|
||||
/// Stored as u32 (integer QPS) for lock-free update; fractional QPS
|
||||
/// information is not useful at the UI grain we're publishing.
|
||||
qps_recent: AtomicU32,
|
||||
failure_count: AtomicU32,
|
||||
failure_window_start: AtomicU64,
|
||||
circuit_open_since: AtomicU64,
|
||||
}
|
||||
|
||||
impl Inference {
|
||||
pub fn new(models: Arc<MLModels>, config: Arc<MLInferenceConfig>) -> Self {
|
||||
let c2_class_idx = Self::find_label_index(&config, "C2 Communication");
|
||||
let normal_class_idx = Self::find_label_index(&config, "Normal");
|
||||
/// Build a new Inference with the given initial state. Use
|
||||
/// `ModelSourceState::Dormant` when no model is loaded (Day 1 default).
|
||||
pub fn new(initial_state: ModelSourceState, config: Arc<MLInferenceConfig>) -> Self {
|
||||
Self {
|
||||
models: ArcSwap::from(models),
|
||||
state: ArcSwap::from_pointee(initial_state),
|
||||
config,
|
||||
c2_class_idx,
|
||||
normal_class_idx,
|
||||
qps_recent: AtomicU32::new(0),
|
||||
failure_count: AtomicU32::new(0),
|
||||
failure_window_start: AtomicU64::new(0),
|
||||
circuit_open_since: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Atomically swap to new models. Old models are dropped when the last reader finishes.
|
||||
pub fn swap_models(&self, new_models: Arc<MLModels>) {
|
||||
self.models.store(new_models);
|
||||
// Reset circuit breaker on successful model swap
|
||||
self.failure_count.store(0, Ordering::Relaxed);
|
||||
self.failure_window_start.store(0, Ordering::Relaxed);
|
||||
self.circuit_open_since.store(0, Ordering::Relaxed);
|
||||
/// Atomically swap in a new state. Any transition INTO `Active` resets
|
||||
/// the circuit breaker so a freshly loaded model starts with a clean
|
||||
/// failure record; `Error ↔ Dormant` transitions leave the CB alone.
|
||||
pub fn swap_state(&self, new_state: ModelSourceState) {
|
||||
let new_is_active = new_state.is_active();
|
||||
self.state.store(Arc::new(new_state));
|
||||
if new_is_active {
|
||||
self.failure_count.store(0, Ordering::Relaxed);
|
||||
self.failure_window_start.store(0, Ordering::Relaxed);
|
||||
self.circuit_open_since.store(0, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
fn find_label_index(config: &MLInferenceConfig, label: &str) -> Option<usize> {
|
||||
config
|
||||
.attack_labels
|
||||
.iter()
|
||||
.find(|(_, v)| v.as_str() == label)
|
||||
.and_then(|(k, _)| k.parse::<usize>().ok())
|
||||
/// Current state snapshot for wire broadcast. Merges the in-memory
|
||||
/// `qps_recent` atomic into the Active info so the UI sees live QPS.
|
||||
pub fn current_status(&self) -> ModelSourceStatus {
|
||||
let guard = self.state.load();
|
||||
let mut status = guard.to_status();
|
||||
if let ModelSourceStatus::Active { ref mut info } = status {
|
||||
info.qps_recent = self.qps_recent.load(Ordering::Relaxed) as f32;
|
||||
}
|
||||
status
|
||||
}
|
||||
|
||||
/// Batched inference with circuit breaker protection.
|
||||
pub fn is_active(&self) -> bool {
|
||||
self.state.load().is_active()
|
||||
}
|
||||
|
||||
/// Per-attack-type confirmations from the active manifest's labels.
|
||||
/// Returns `None` when the source is Dormant/Error or when the label has
|
||||
/// no `confirmations` override; callers apply their own fallback.
|
||||
pub fn confirmations_for_attack_type(&self, attack_type_name: &str) -> Option<usize> {
|
||||
match self.state.load().as_ref() {
|
||||
ModelSourceState::Active { adapter, .. } => adapter.confirmations_for(attack_type_name),
|
||||
ModelSourceState::Dormant | ModelSourceState::Error { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Batched inference with circuit breaker protection + state dispatch.
|
||||
/// Returns an empty Vec when Dormant / Error / circuit-open.
|
||||
pub fn infer_batch(&self, flows: &[FlowData]) -> Vec<DetectionResult> {
|
||||
if self.is_circuit_open() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
match panic::catch_unwind(AssertUnwindSafe(|| self.infer_batch_inner(flows))) {
|
||||
// Snapshot the state once so the whole batch sees a consistent adapter.
|
||||
let guard = self.state.load();
|
||||
let adapter = match guard.as_ref() {
|
||||
ModelSourceState::Active { adapter, .. } => adapter,
|
||||
ModelSourceState::Dormant | ModelSourceState::Error { .. } => {
|
||||
return Vec::new();
|
||||
}
|
||||
};
|
||||
|
||||
match panic::catch_unwind(AssertUnwindSafe(|| self.infer_batch_inner(adapter, flows))) {
|
||||
Ok(results) => {
|
||||
// Success: reset failure counter
|
||||
self.failure_count.store(0, Ordering::Relaxed);
|
||||
self.failure_window_start.store(0, Ordering::Relaxed);
|
||||
results
|
||||
@ -91,23 +136,68 @@ impl Inference {
|
||||
}
|
||||
}
|
||||
|
||||
fn infer_batch_inner(&self, flows: &[FlowData]) -> Vec<DetectionResult> {
|
||||
let models = self.models.load();
|
||||
let n = flows.len();
|
||||
let batch_size = models.batch_size;
|
||||
let n_ae = self.config.num_ae_features();
|
||||
let n_cls = self.config.num_classifier_features();
|
||||
/// Dispatch on adapter variant.
|
||||
fn infer_batch_inner(&self, adapter: &MLModelAdapter, flows: &[FlowData]) -> Vec<DetectionResult> {
|
||||
match adapter {
|
||||
MLModelAdapter::MultiTask {
|
||||
ae,
|
||||
classifier,
|
||||
batch_size,
|
||||
n_ae,
|
||||
n_cls,
|
||||
labels,
|
||||
normal_idx,
|
||||
c2_idx,
|
||||
} => self.infer_multitask(
|
||||
ae,
|
||||
classifier,
|
||||
*batch_size,
|
||||
*n_ae,
|
||||
*n_cls,
|
||||
labels,
|
||||
*normal_idx,
|
||||
*c2_idx,
|
||||
flows,
|
||||
),
|
||||
MLModelAdapter::AutoencoderOnly {
|
||||
model,
|
||||
batch_size,
|
||||
n_features,
|
||||
} => self.infer_autoencoder_only(model, *batch_size, *n_features, flows),
|
||||
MLModelAdapter::ClassifierOnly {
|
||||
model,
|
||||
batch_size,
|
||||
n_features,
|
||||
labels,
|
||||
normal_idx,
|
||||
} => self.infer_classifier_only(model, *batch_size, *n_features, labels, *normal_idx, flows),
|
||||
}
|
||||
}
|
||||
|
||||
/// MultiTask path. Runs the AE batch → computes per-flow MSE → feeds the
|
||||
/// classifier over (ae_features ++ ae_score) → fires on anomaly OR
|
||||
/// non-Normal classifier agreement OR elevated C2 head.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn infer_multitask(
|
||||
&self,
|
||||
ae: &RunnableModel,
|
||||
classifier: &RunnableModel,
|
||||
batch_size: usize,
|
||||
n_ae: usize,
|
||||
n_cls: usize,
|
||||
labels: &BTreeMap<String, LabelSpec>,
|
||||
normal_idx: Option<usize>,
|
||||
c2_idx: Option<usize>,
|
||||
flows: &[FlowData],
|
||||
) -> Vec<DetectionResult> {
|
||||
let n = flows.len();
|
||||
|
||||
// Phase 1: preprocess all AE features
|
||||
let all_ae_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
|
||||
|
||||
// Phase 2: run AE in chunks → compute per-flow MSE
|
||||
let mut ae_scores = Vec::with_capacity(n);
|
||||
for chunk_start in (0..n).step_by(batch_size) {
|
||||
let chunk_end = (chunk_start + batch_size).min(n);
|
||||
let actual = chunk_end - chunk_start;
|
||||
|
||||
// Build (batch_size, n_ae) tensor, zero-padded
|
||||
let ae_input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_ae), |(i, j)| {
|
||||
if i < actual {
|
||||
all_ae_features[chunk_start + i][j]
|
||||
@ -116,7 +206,7 @@ impl Inference {
|
||||
}
|
||||
});
|
||||
|
||||
match Self::run_ae_batch(&models, &ae_input, actual, n_ae) {
|
||||
match run_ae_batch(ae, &ae_input, actual, n_ae) {
|
||||
Ok(scores) => ae_scores.extend_from_slice(&scores),
|
||||
Err(e) => {
|
||||
log!(MLLog::InferenceFailed("DeepAutoEncoder".to_string(), e.to_string()));
|
||||
@ -126,7 +216,6 @@ impl Inference {
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 3: build classifier input (ae_features ++ ae_score) and run in chunks
|
||||
let mut all_anomaly = Vec::with_capacity(n);
|
||||
let mut all_class_probs = Vec::with_capacity(n);
|
||||
let mut all_c2_scores = Vec::with_capacity(n);
|
||||
@ -135,7 +224,6 @@ impl Inference {
|
||||
let chunk_end = (chunk_start + batch_size).min(n);
|
||||
let actual = chunk_end - chunk_start;
|
||||
|
||||
// Build (batch_size, n_cls) tensor
|
||||
let cls_input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_cls), |(i, j)| {
|
||||
if i < actual {
|
||||
if j < n_ae {
|
||||
@ -148,7 +236,7 @@ impl Inference {
|
||||
}
|
||||
});
|
||||
|
||||
match Self::run_classifier_batch(&models, &cls_input, actual) {
|
||||
match run_classifier_batch(classifier, &cls_input, actual) {
|
||||
Ok((anomaly, class_probs, c2)) => {
|
||||
all_anomaly.extend_from_slice(&anomaly);
|
||||
all_class_probs.extend(class_probs);
|
||||
@ -162,159 +250,192 @@ impl Inference {
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 4: build DetectionResults
|
||||
let mut results = Vec::with_capacity(n);
|
||||
let class_min_conf = self.config.class_min_confidence;
|
||||
let anomaly_thr = self.config.anomaly_threshold;
|
||||
let c2_thr = self.config.c2_threshold;
|
||||
|
||||
for i in 0..n {
|
||||
let flow = &flows[i];
|
||||
let flow_key = format!(
|
||||
"{}:{} -> {}:{} (proto {}) [{}]",
|
||||
flow.flow_key.src_ip_string(),
|
||||
flow.flow_key.src_port,
|
||||
flow.flow_key.dst_ip_string(),
|
||||
flow.flow_key.dst_port,
|
||||
flow.flow_key.protocol,
|
||||
flow.direction
|
||||
);
|
||||
let class_probs = &all_class_probs[i];
|
||||
let anomaly = all_anomaly[i];
|
||||
let c2 = all_c2_scores[i];
|
||||
|
||||
let result = self.build_detection_result(
|
||||
flow,
|
||||
flow_key,
|
||||
ae_scores[i],
|
||||
all_anomaly[i],
|
||||
&all_class_probs[i],
|
||||
all_c2_scores[i],
|
||||
);
|
||||
results.push(result);
|
||||
let (predicted_class, class_confidence) = argmax(class_probs);
|
||||
let attack_type_name = labels
|
||||
.get(&predicted_class.to_string())
|
||||
.map(|l| l.name.clone())
|
||||
.unwrap_or_else(|| "UNKNOWN".to_string());
|
||||
|
||||
let classifier_fires =
|
||||
normal_idx.is_none_or(|ni| predicted_class != ni) && class_confidence >= class_min_conf;
|
||||
let c2_fires = c2 > c2_thr;
|
||||
let mut is_attack = anomaly > anomaly_thr || classifier_fires || c2_fires;
|
||||
|
||||
// "Normal" with no C2 elevation stays benign regardless of AE noise.
|
||||
if normal_idx == Some(predicted_class) && !c2_fires {
|
||||
is_attack = false;
|
||||
}
|
||||
|
||||
// When the manifest declares a C2 class and its head score beats
|
||||
// the classifier's probability for that same class, relabel the
|
||||
// event with the manifest's C2 label and use the head score as
|
||||
// the outgoing confidence. Manifests without a C2 class keep the
|
||||
// argmax label and class-confidence untouched.
|
||||
let mut attack_type = attack_type_name;
|
||||
let mut confidence = class_confidence;
|
||||
if c2_fires
|
||||
&& let Some(idx) = c2_idx
|
||||
&& let Some(c2_class_prob) = class_probs.get(idx).copied()
|
||||
&& c2 > c2_class_prob
|
||||
{
|
||||
if let Some(spec) = labels.get(&idx.to_string()) {
|
||||
attack_type = spec.name.clone();
|
||||
}
|
||||
confidence = c2;
|
||||
}
|
||||
|
||||
results.push(DetectionResult {
|
||||
flow_key: build_flow_key_label(flow),
|
||||
flow_key_raw: flow.flow_key.clone(),
|
||||
direction: flow.direction,
|
||||
is_attack,
|
||||
attack_type: if is_attack { Some(attack_type) } else { None },
|
||||
confidence,
|
||||
ae_score: ae_scores[i],
|
||||
anomaly_score: anomaly,
|
||||
c2_score: c2,
|
||||
packet_count: flow.packet_count() as u64,
|
||||
flow_duration_us: flow.duration_us(),
|
||||
});
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Run AE on a padded batch, return MSE scores for the first `actual` rows.
|
||||
fn run_ae_batch(
|
||||
models: &MLModels,
|
||||
input: &tract_ndarray::Array2<f32>,
|
||||
actual: usize,
|
||||
n_features: usize,
|
||||
) -> TractResult<Vec<f32>> {
|
||||
let result = models.deep_autoencoder.run(tvec![input.clone().into_tensor().into()])?;
|
||||
|
||||
let output = result[0]
|
||||
.to_array_view::<f32>()?
|
||||
.into_dimensionality::<tract_ndarray::Ix2>()?;
|
||||
|
||||
let diff = input - &output;
|
||||
let sq = &diff * &diff;
|
||||
|
||||
let mut scores = Vec::with_capacity(actual);
|
||||
let n_f = n_features as f32;
|
||||
for i in 0..actual {
|
||||
let mse: f32 = sq.row(i).sum() / n_f;
|
||||
scores.push(mse);
|
||||
}
|
||||
Ok(scores)
|
||||
}
|
||||
|
||||
/// Returns (anomaly_scores, per_class_probs, c2_scores) for the first `actual` rows.
|
||||
fn run_classifier_batch(
|
||||
models: &MLModels,
|
||||
input: &tract_ndarray::Array2<f32>,
|
||||
actual: usize,
|
||||
) -> TractResult<ClassifierBatchOutput> {
|
||||
let result = models.classifier.run(tvec![input.clone().into_tensor().into()])?;
|
||||
|
||||
// Output 0: anomaly (batch_size, 1)
|
||||
let anomaly_view = result[0].to_array_view::<f32>()?;
|
||||
let anomaly: Vec<f32> = (0..actual)
|
||||
.map(|i| anomaly_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
|
||||
.collect();
|
||||
|
||||
// Output 1: class_probs (batch_size, n_classes)
|
||||
let class_view = result[1]
|
||||
.to_array_view::<f32>()?
|
||||
.into_dimensionality::<tract_ndarray::Ix2>()?;
|
||||
let class_probs: Vec<Vec<f32>> = (0..actual)
|
||||
.map(|i| class_view.row(i).iter().copied().collect())
|
||||
.collect();
|
||||
|
||||
// Output 2: c2_score (batch_size, 1)
|
||||
let c2_view = result[2].to_array_view::<f32>()?;
|
||||
let c2: Vec<f32> = (0..actual)
|
||||
.map(|i| c2_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
|
||||
.collect();
|
||||
|
||||
Ok((anomaly, class_probs, c2))
|
||||
}
|
||||
|
||||
fn build_detection_result(
|
||||
/// AutoencoderOnly path. Output is reconstruction MSE; when it exceeds
|
||||
/// `anomaly_threshold` the flow is marked as a generic `anomaly`. There
|
||||
/// are no classifier outputs, so no per-class gating happens here.
|
||||
fn infer_autoencoder_only(
|
||||
&self,
|
||||
flow: &FlowData,
|
||||
flow_key: String,
|
||||
ae_score: f32,
|
||||
anomaly_score: f32,
|
||||
class_probs: &[f32],
|
||||
c2_score: f32,
|
||||
) -> DetectionResult {
|
||||
let is_attack = anomaly_score > self.config.anomaly_threshold;
|
||||
model: &RunnableModel,
|
||||
batch_size: usize,
|
||||
n_features: usize,
|
||||
flows: &[FlowData],
|
||||
) -> Vec<DetectionResult> {
|
||||
let n = flows.len();
|
||||
let all_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
|
||||
let mut scores = Vec::with_capacity(n);
|
||||
|
||||
let (predicted_class, class_confidence) = class_probs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(CmpOrdering::Equal))
|
||||
.map(|(i, &p)| (i, p))
|
||||
.unwrap_or((0, 0.0));
|
||||
|
||||
let mut attack_type = self
|
||||
.config
|
||||
.attack_labels
|
||||
.get(&predicted_class.to_string())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "UNKNOWN".to_string());
|
||||
|
||||
let mut confidence = class_confidence;
|
||||
|
||||
// C2 head override
|
||||
if c2_score > self.config.c2_threshold {
|
||||
let c2_class_prob = self
|
||||
.c2_class_idx
|
||||
.and_then(|idx| class_probs.get(idx).copied())
|
||||
.unwrap_or(0.0);
|
||||
if c2_score > c2_class_prob {
|
||||
attack_type = "C2 Communication".to_string();
|
||||
confidence = c2_score;
|
||||
for chunk_start in (0..n).step_by(batch_size) {
|
||||
let chunk_end = (chunk_start + batch_size).min(n);
|
||||
let actual = chunk_end - chunk_start;
|
||||
let input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_features), |(i, j)| {
|
||||
if i < actual {
|
||||
all_features[chunk_start + i][j]
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
});
|
||||
match run_ae_batch(model, &input, actual, n_features) {
|
||||
Ok(s) => scores.extend_from_slice(&s),
|
||||
Err(e) => {
|
||||
log!(MLLog::InferenceFailed("AutoencoderOnly".to_string(), e.to_string()));
|
||||
self.record_failure();
|
||||
return Vec::new();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// "Normal" class = benign
|
||||
if Some(predicted_class) == self.normal_class_idx && c2_score <= self.config.c2_threshold {
|
||||
return DetectionResult {
|
||||
flow_key,
|
||||
flow_key_raw: flow.flow_key.clone(),
|
||||
direction: flow.direction,
|
||||
is_attack: false,
|
||||
attack_type: None,
|
||||
confidence: class_confidence,
|
||||
ae_score,
|
||||
anomaly_score,
|
||||
c2_score,
|
||||
packet_count: flow.packet_count() as u64,
|
||||
flow_duration_us: flow.duration_us(),
|
||||
let thr = self.config.anomaly_threshold;
|
||||
flows
|
||||
.iter()
|
||||
.zip(scores.iter())
|
||||
.map(|(flow, &score)| {
|
||||
let is_attack = score > thr;
|
||||
DetectionResult {
|
||||
flow_key: build_flow_key_label(flow),
|
||||
flow_key_raw: flow.flow_key.clone(),
|
||||
direction: flow.direction,
|
||||
is_attack,
|
||||
attack_type: if is_attack { Some("anomaly".to_string()) } else { None },
|
||||
// AE-only has no separate classifier confidence; reuse the score.
|
||||
confidence: score,
|
||||
ae_score: score,
|
||||
anomaly_score: score,
|
||||
c2_score: 0.0,
|
||||
packet_count: flow.packet_count() as u64,
|
||||
flow_duration_us: flow.duration_us(),
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// ClassifierOnly path. Output is per-class softmax; the manifest's
|
||||
/// labels drive attack_type and the flow fires when the argmax class
|
||||
/// isn't Normal and confidence ≥ `class_min_confidence`.
|
||||
fn infer_classifier_only(
|
||||
&self,
|
||||
model: &RunnableModel,
|
||||
batch_size: usize,
|
||||
n_features: usize,
|
||||
labels: &BTreeMap<String, LabelSpec>,
|
||||
normal_idx: Option<usize>,
|
||||
flows: &[FlowData],
|
||||
) -> Vec<DetectionResult> {
|
||||
let n = flows.len();
|
||||
let all_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
|
||||
let class_min_conf = self.config.class_min_confidence;
|
||||
let mut results = Vec::with_capacity(n);
|
||||
|
||||
for chunk_start in (0..n).step_by(batch_size) {
|
||||
let chunk_end = (chunk_start + batch_size).min(n);
|
||||
let actual = chunk_end - chunk_start;
|
||||
let input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_features), |(i, j)| {
|
||||
if i < actual {
|
||||
all_features[chunk_start + i][j]
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
});
|
||||
let class_probs = match run_classifier_only_batch(model, &input, actual) {
|
||||
Ok(cp) => cp,
|
||||
Err(e) => {
|
||||
log!(MLLog::InferenceFailed("ClassifierOnly".to_string(), e.to_string()));
|
||||
self.record_failure();
|
||||
return Vec::new();
|
||||
}
|
||||
};
|
||||
|
||||
for (i, probs) in class_probs.into_iter().enumerate() {
|
||||
let flow = &flows[chunk_start + i];
|
||||
let (predicted_class, confidence) = argmax(&probs);
|
||||
let is_attack = normal_idx.is_none_or(|ni| predicted_class != ni) && confidence >= class_min_conf;
|
||||
let attack_type = if is_attack {
|
||||
labels
|
||||
.get(&predicted_class.to_string())
|
||||
.map(|l| l.name.clone())
|
||||
.unwrap_or_else(|| "UNKNOWN".to_string())
|
||||
} else {
|
||||
"Normal".to_string()
|
||||
};
|
||||
results.push(DetectionResult {
|
||||
flow_key: build_flow_key_label(flow),
|
||||
flow_key_raw: flow.flow_key.clone(),
|
||||
direction: flow.direction,
|
||||
is_attack,
|
||||
attack_type: if is_attack { Some(attack_type) } else { None },
|
||||
confidence,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
packet_count: flow.packet_count() as u64,
|
||||
flow_duration_us: flow.duration_us(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
DetectionResult {
|
||||
flow_key,
|
||||
flow_key_raw: flow.flow_key.clone(),
|
||||
direction: flow.direction,
|
||||
is_attack,
|
||||
attack_type: if is_attack { Some(attack_type) } else { None },
|
||||
confidence,
|
||||
ae_score,
|
||||
anomaly_score,
|
||||
c2_score,
|
||||
packet_count: flow.packet_count() as u64,
|
||||
flow_duration_us: flow.duration_us(),
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
fn preprocess_ae_features(&self, flow: &FlowData) -> Vec<f32> {
|
||||
@ -325,6 +446,13 @@ impl Inference {
|
||||
features.features.iter().map(|&x| x as f32).collect()
|
||||
}
|
||||
|
||||
/// Publish a rolling QPS estimate visible in `current_status().info.qps_recent`.
|
||||
/// Called by the engine after each inference tick completes.
|
||||
pub fn record_tick_qps(&self, flows_per_second: f32) {
|
||||
self.qps_recent
|
||||
.store(flows_per_second.max(0.0) as u32, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// -- Circuit breaker ---------------------------------------------------------
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
@ -341,7 +469,6 @@ impl Inference {
|
||||
}
|
||||
let elapsed = Self::now_secs().saturating_sub(open_since);
|
||||
if elapsed >= CIRCUIT_BREAKER_COOLDOWN_SECS {
|
||||
// Cooldown elapsed — reset and allow inference
|
||||
self.circuit_open_since.store(0, Ordering::Relaxed);
|
||||
self.failure_count.store(0, Ordering::Relaxed);
|
||||
self.failure_window_start.store(0, Ordering::Relaxed);
|
||||
@ -355,7 +482,6 @@ impl Inference {
|
||||
let now = Self::now_secs();
|
||||
let window_start = self.failure_window_start.load(Ordering::Relaxed);
|
||||
|
||||
// If window has expired, start a new window
|
||||
if window_start == 0 || now.saturating_sub(window_start) > CIRCUIT_BREAKER_WINDOW_SECS {
|
||||
self.failure_window_start.store(now, Ordering::Relaxed);
|
||||
self.failure_count.store(1, Ordering::Relaxed);
|
||||
@ -369,3 +495,117 @@ impl Inference {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- Free helpers (can be unit-tested without an Inference) -------------------
|
||||
|
||||
fn argmax(probs: &[f32]) -> (usize, f32) {
|
||||
probs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(CmpOrdering::Equal))
|
||||
.map(|(i, &p)| (i, p))
|
||||
.unwrap_or((0, 0.0))
|
||||
}
|
||||
|
||||
fn build_flow_key_label(flow: &FlowData) -> String {
|
||||
format!(
|
||||
"{}:{} -> {}:{} (proto {}) [{}]",
|
||||
flow.flow_key.src_ip_string(),
|
||||
flow.flow_key.src_port,
|
||||
flow.flow_key.dst_ip_string(),
|
||||
flow.flow_key.dst_port,
|
||||
flow.flow_key.protocol,
|
||||
flow.direction
|
||||
)
|
||||
}
|
||||
|
||||
/// Run an autoencoder-style model: output shape == input shape, score is
|
||||
/// per-row mean squared error between input and reconstruction.
|
||||
fn run_ae_batch(
|
||||
model: &RunnableModel,
|
||||
input: &tract_ndarray::Array2<f32>,
|
||||
actual: usize,
|
||||
n_features: usize,
|
||||
) -> TractResult<Vec<f32>> {
|
||||
let result = model.run(tvec![input.clone().into_tensor().into()])?;
|
||||
let output = result[0]
|
||||
.to_array_view::<f32>()?
|
||||
.into_dimensionality::<tract_ndarray::Ix2>()?;
|
||||
let diff = input - &output;
|
||||
let sq = &diff * &diff;
|
||||
|
||||
let mut scores = Vec::with_capacity(actual);
|
||||
let n_f = n_features as f32;
|
||||
for i in 0..actual {
|
||||
scores.push(sq.row(i).sum() / n_f);
|
||||
}
|
||||
Ok(scores)
|
||||
}
|
||||
|
||||
/// Run a 3-output multi-task classifier: (anomaly, class_probs, c2_score).
|
||||
fn run_classifier_batch(
|
||||
model: &RunnableModel,
|
||||
input: &tract_ndarray::Array2<f32>,
|
||||
actual: usize,
|
||||
) -> TractResult<ClassifierBatchOutput> {
|
||||
let result = model.run(tvec![input.clone().into_tensor().into()])?;
|
||||
|
||||
let anomaly_view = result[0].to_array_view::<f32>()?;
|
||||
let anomaly: Vec<f32> = (0..actual)
|
||||
.map(|i| anomaly_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
|
||||
.collect();
|
||||
|
||||
let class_view = result[1]
|
||||
.to_array_view::<f32>()?
|
||||
.into_dimensionality::<tract_ndarray::Ix2>()?;
|
||||
let class_probs: Vec<Vec<f32>> = (0..actual)
|
||||
.map(|i| class_view.row(i).iter().copied().collect())
|
||||
.collect();
|
||||
|
||||
let c2_view = result[2].to_array_view::<f32>()?;
|
||||
let c2: Vec<f32> = (0..actual)
|
||||
.map(|i| c2_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
|
||||
.collect();
|
||||
|
||||
Ok((anomaly, class_probs, c2))
|
||||
}
|
||||
|
||||
/// Run a single-output classifier (ClassifierOnly adapter).
|
||||
fn run_classifier_only_batch(
|
||||
model: &RunnableModel,
|
||||
input: &tract_ndarray::Array2<f32>,
|
||||
actual: usize,
|
||||
) -> TractResult<Vec<Vec<f32>>> {
|
||||
let result = model.run(tvec![input.clone().into_tensor().into()])?;
|
||||
let view = result[0]
|
||||
.to_array_view::<f32>()?
|
||||
.into_dimensionality::<tract_ndarray::Ix2>()?;
|
||||
Ok((0..actual).map(|i| view.row(i).iter().copied().collect()).collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn argmax_picks_highest() {
|
||||
let (idx, val) = argmax(&[0.1, 0.5, 0.3, 0.4]);
|
||||
assert_eq!(idx, 1);
|
||||
assert!((val - 0.5).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn argmax_empty_returns_zero() {
|
||||
assert_eq!(argmax(&[]), (0, 0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn argmax_equal_picks_last() {
|
||||
// Iterator::max_by returns the LAST element when comparisons are
|
||||
// equal (in contrast to min_by). Ties in softmax probabilities are
|
||||
// rare in practice, and "last wins" is a consistent contract across
|
||||
// this codebase.
|
||||
let (idx, _) = argmax(&[0.25, 0.25, 0.25, 0.25]);
|
||||
assert_eq!(idx, 3);
|
||||
}
|
||||
}
|
||||
|
||||
@ -39,8 +39,7 @@ pub struct ModelPaths {
|
||||
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.
|
||||
/// Per-label metadata driving classifier output decoding and alert gating.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct LabelSpec {
|
||||
pub name: String,
|
||||
@ -60,10 +59,11 @@ pub struct Thresholds {
|
||||
pub class_min_confidence: Option<f32>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub ae: Option<f32>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub alert_multiplier: Option<f32>,
|
||||
}
|
||||
|
||||
/// Pointer to a preprocessing sidecar (scaler/clip arrays). For v10 this is
|
||||
/// the legacy `inference_config.json`.
|
||||
/// Pointer to the JSON sidecar carrying scaler / clip / weight arrays.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Preprocessing {
|
||||
pub scaler_sidecar: String,
|
||||
@ -134,6 +134,43 @@ impl ModelManifest {
|
||||
}
|
||||
}
|
||||
}
|
||||
self.validate_labels(path)?;
|
||||
self.validate_thresholds(path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_labels(&self, path: &Path) -> Result<(), MLError> {
|
||||
let mut seen: Vec<String> = Vec::with_capacity(self.labels.len());
|
||||
for spec in self.labels.values() {
|
||||
if let Some(n) = spec.confirmations
|
||||
&& n == 0
|
||||
{
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
format!("label '{}' has confirmations: 0 (must be ≥ 1)", spec.name),
|
||||
));
|
||||
}
|
||||
let lower = spec.name.to_ascii_lowercase();
|
||||
if seen.iter().any(|s| s == &lower) {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
format!("duplicate label name '{}' (case-insensitive)", spec.name),
|
||||
));
|
||||
}
|
||||
seen.push(lower);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_thresholds(&self, path: &Path) -> Result<(), MLError> {
|
||||
if let Some(m) = self.thresholds.alert_multiplier
|
||||
&& !(m.is_finite() && m > 0.0)
|
||||
{
|
||||
return Err(MLError::ManifestInvalid(
|
||||
path.to_path_buf(),
|
||||
format!("thresholds.alert_multiplier must be finite and > 0 (got {m})"),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -226,4 +263,59 @@ features: []
|
||||
let err = parsed.validate(path).expect_err("should reject empty features");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_confirmations() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: classifier_only
|
||||
models:
|
||||
model: m.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
labels:
|
||||
"0": { name: Bot, confirmations: 0 }
|
||||
"#;
|
||||
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 confirmations: 0");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_duplicate_label_name_ignoring_case() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: classifier_only
|
||||
models:
|
||||
model: m.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
labels:
|
||||
"0": { name: Bot }
|
||||
"1": { name: BOT }
|
||||
"#;
|
||||
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 duplicate label names");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_positive_alert_multiplier() {
|
||||
let yaml = r#"
|
||||
name: bad
|
||||
adapter: classifier_only
|
||||
models:
|
||||
model: m.onnx
|
||||
features:
|
||||
- flow_duration
|
||||
thresholds:
|
||||
alert_multiplier: 0
|
||||
"#;
|
||||
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 alert_multiplier <= 0");
|
||||
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
pub mod adapter;
|
||||
pub mod aggregator;
|
||||
pub mod alert;
|
||||
pub mod config_loader;
|
||||
|
||||
@ -1,4 +1,13 @@
|
||||
//! Factory that builds an `MLModelAdapter` from a `ModelManifest`: inspects
|
||||
//! the manifest's `adapter` field, loads the named ONNX file(s) with shape
|
||||
//! validation against the inference config's feature counts, and wraps the
|
||||
//! `RunnableModel`s in `Arc` so hot-reload can swap without per-tick clones.
|
||||
//! A shape mismatch yields `MLError::FeatureMismatch` carrying both
|
||||
//! expected and observed dims for the upload UI to render.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use macros::log;
|
||||
@ -6,105 +15,153 @@ 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 super::adapter::MLModelAdapter;
|
||||
use super::manifest::{AdapterKind, LabelSpec, ModelManifest};
|
||||
use crate::model::config::constants::MODELS_DIR;
|
||||
use crate::model::detection::ml_detection::RunnableModel;
|
||||
use crate::model::error::ml::MLError;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
|
||||
pub struct MLModels {
|
||||
pub deep_autoencoder: RunnableModel,
|
||||
pub classifier: RunnableModel,
|
||||
pub batch_size: usize,
|
||||
/// Build an `MLModelAdapter` by loading the ONNX file(s) the manifest names,
|
||||
/// validating shape against the inference config's feature counts, and
|
||||
/// wrapping the underlying `RunnableModel`s in `Arc` for zero-copy swap.
|
||||
pub fn build_adapter(
|
||||
manifest: &ModelManifest,
|
||||
manifest_path: Option<&Path>,
|
||||
inference_config: &MLInferenceConfig,
|
||||
batch_size: usize,
|
||||
) -> Result<MLModelAdapter, MLError> {
|
||||
let resolve = |rel: &str| -> PathBuf {
|
||||
match manifest_path {
|
||||
Some(mp) => ModelManifest::resolve_relative(mp, rel),
|
||||
None => PathBuf::from(MODELS_DIR).join(rel),
|
||||
}
|
||||
};
|
||||
|
||||
match manifest.adapter {
|
||||
AdapterKind::AutoencoderOnly => {
|
||||
let Some(model_name) = manifest.models.model.as_deref() else {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
|
||||
"autoencoder_only adapter requires models.model".to_string(),
|
||||
));
|
||||
};
|
||||
let path = resolve(model_name);
|
||||
let n_features = inference_config.num_ae_features();
|
||||
let model = Arc::new(loader(&path, model_name, n_features, batch_size)?);
|
||||
Ok(MLModelAdapter::AutoencoderOnly {
|
||||
model,
|
||||
batch_size,
|
||||
n_features,
|
||||
})
|
||||
}
|
||||
AdapterKind::ClassifierOnly => {
|
||||
let Some(model_name) = manifest.models.model.as_deref() else {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
|
||||
"classifier_only adapter requires models.model".to_string(),
|
||||
));
|
||||
};
|
||||
let path = resolve(model_name);
|
||||
let n_features = inference_config.num_classifier_features();
|
||||
let model = Arc::new(loader(&path, model_name, n_features, batch_size)?);
|
||||
let labels = manifest.labels.clone();
|
||||
let normal_idx = find_label_index(&labels, "Normal");
|
||||
Ok(MLModelAdapter::ClassifierOnly {
|
||||
model,
|
||||
batch_size,
|
||||
n_features,
|
||||
labels,
|
||||
normal_idx,
|
||||
})
|
||||
}
|
||||
AdapterKind::MultiTask => {
|
||||
let Some(ae_name) = manifest.models.autoencoder.as_deref() else {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
|
||||
"multi_task adapter requires models.autoencoder".to_string(),
|
||||
));
|
||||
};
|
||||
let Some(cls_name) = manifest.models.classifier.as_deref() else {
|
||||
return Err(MLError::ManifestInvalid(
|
||||
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
|
||||
"multi_task adapter requires models.classifier".to_string(),
|
||||
));
|
||||
};
|
||||
let n_ae = inference_config.num_ae_features();
|
||||
let n_cls = inference_config.num_classifier_features();
|
||||
let ae = Arc::new(loader(&resolve(ae_name), ae_name, n_ae, batch_size)?);
|
||||
let classifier = Arc::new(loader(&resolve(cls_name), cls_name, n_cls, batch_size)?);
|
||||
let labels = manifest.labels.clone();
|
||||
let normal_idx = find_label_index(&labels, "Normal");
|
||||
let c2_idx = find_label_index(&labels, "C2 Communication");
|
||||
Ok(MLModelAdapter::MultiTask {
|
||||
ae,
|
||||
classifier,
|
||||
batch_size,
|
||||
n_ae,
|
||||
n_cls,
|
||||
labels,
|
||||
normal_idx,
|
||||
c2_idx,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MLModels {
|
||||
pub fn load_models(
|
||||
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,
|
||||
)
|
||||
}
|
||||
/// Locate a label index by its `name` field. Used to cache hot-path indices
|
||||
/// (Normal, C2 Communication) at adapter build time rather than re-scanning
|
||||
/// the label map on every inference batch.
|
||||
fn find_label_index(labels: &BTreeMap<String, LabelSpec>, target: &str) -> Option<usize> {
|
||||
labels
|
||||
.iter()
|
||||
.find(|(_, v)| v.name.eq_ignore_ascii_case(target))
|
||||
.and_then(|(k, _)| k.parse::<usize>().ok())
|
||||
}
|
||||
|
||||
/// 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(autoencoder_name, inference_config.num_ae_features(), batch_size)?,
|
||||
classifier: Self::loader(classifier_name, inference_config.num_classifier_features(), batch_size)?,
|
||||
batch_size,
|
||||
})
|
||||
}
|
||||
fn loader(model_path: &Path, model_name: &str, features: usize, batch_size: usize) -> Result<RunnableModel, MLError> {
|
||||
log!(MLLog::ModelLoading(model_name.to_string(), features, batch_size));
|
||||
|
||||
fn loader(model: &str, features: usize, batch_size: usize) -> Result<RunnableModel, MLError> {
|
||||
let model_path = PathBuf::from("models").join(model);
|
||||
let start = Instant::now();
|
||||
let result = loader_inner(model_path, model_name, features, batch_size);
|
||||
let elapsed_ms = start.elapsed().as_millis() as u64;
|
||||
log!(MLLog::ModelLoadComplete(model_name.to_string(), elapsed_ms));
|
||||
|
||||
log!(MLLog::ModelLoading(model.to_string(), features, batch_size));
|
||||
result
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
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));
|
||||
fn loader_inner(
|
||||
model_path: &Path,
|
||||
model_name: &str,
|
||||
features: usize,
|
||||
batch_size: usize,
|
||||
) -> Result<RunnableModel, MLError> {
|
||||
let mut onnx_model = onnx()
|
||||
.model_for_path(model_path)
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("parse ONNX: {e}")))?;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn loader_inner(
|
||||
model_path: &Path,
|
||||
model: &str,
|
||||
features: usize,
|
||||
batch_size: usize,
|
||||
) -> Result<RunnableModel, MLError> {
|
||||
let mut onnx_model = onnx()
|
||||
.model_for_path(model_path)
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("parse ONNX: {e}")))?;
|
||||
|
||||
// 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.to_path_buf(), features, onnx_dim));
|
||||
}
|
||||
if let Some(onnx_dim) = introspect_input_features(&onnx_model) {
|
||||
let matched = onnx_dim == features;
|
||||
log!(MLLog::OnnxShapeChecked(
|
||||
model_name.to_string(),
|
||||
features,
|
||||
onnx_dim,
|
||||
matched,
|
||||
));
|
||||
if !matched {
|
||||
return Err(MLError::FeatureMismatch(model_path.to_path_buf(), features, onnx_dim));
|
||||
}
|
||||
|
||||
onnx_model
|
||||
.set_input_fact(0, f32::fact([batch_size, features]).into())
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("set_input_fact: {e}")))?;
|
||||
|
||||
onnx_model
|
||||
.into_optimized()
|
||||
.and_then(|m| m.into_runnable())
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("optimize/runnable: {e}")))
|
||||
}
|
||||
|
||||
pub fn get_model_info(&self, name: &str) -> String {
|
||||
let model = match name {
|
||||
"deep_autoencoder" => &self.deep_autoencoder,
|
||||
"classifier" => &self.classifier,
|
||||
_ => return "unknown model".to_string(),
|
||||
};
|
||||
onnx_model
|
||||
.set_input_fact(0, f32::fact([batch_size, features]).into())
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("set_input_fact: {e}")))?;
|
||||
|
||||
let inputs = model.model().inputs.len();
|
||||
let outputs = model.model().outputs.len();
|
||||
format!(
|
||||
"{name}: inputs: {inputs}, outputs: {outputs}, batch_size: {}",
|
||||
self.batch_size
|
||||
)
|
||||
}
|
||||
onnx_model
|
||||
.into_optimized()
|
||||
.and_then(|m| m.into_runnable())
|
||||
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("optimize/runnable: {e}")))
|
||||
}
|
||||
|
||||
/// Read the concrete last-dim (feature count) from an ONNX model's declared input fact.
|
||||
@ -123,8 +180,8 @@ fn introspect_input_features(model: &InferenceModel) -> Option<usize> {
|
||||
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).
|
||||
/// Repo-root integration: the shipped v10 AE ONNX's input last-dim must
|
||||
/// still be 31. A change here means the shipped manifest + sidecar drifted.
|
||||
#[test]
|
||||
fn introspect_v10_autoencoder_input_is_31() {
|
||||
let ae_path = PathBuf::from("models/deep_autoencoder.onnx");
|
||||
@ -149,24 +206,23 @@ mod tests {
|
||||
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.
|
||||
/// Smoke: the shipped manifest + sidecar must load and produce a
|
||||
/// `MultiTask` adapter with both underlying models.
|
||||
#[test]
|
||||
fn v10_load_named_with_manifest_paths_succeeds() {
|
||||
fn v10_build_adapter_multitask() {
|
||||
let manifest_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");
|
||||
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(manifest_path).expect("config load");
|
||||
let adapter = build_adapter(&manifest, Some(manifest_path), &cfg, 8).expect("build adapter");
|
||||
match adapter {
|
||||
MLModelAdapter::MultiTask { n_ae, n_cls, .. } => {
|
||||
assert_eq!(n_ae, 31);
|
||||
assert_eq!(n_cls, 32);
|
||||
}
|
||||
_ => panic!("v10 manifest should build MultiTask adapter"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,43 +1,41 @@
|
||||
use std::path::PathBuf;
|
||||
//! Filesystem watcher over `models/`. Reloads the ML source whenever the
|
||||
//! manifest or ONNX files change, re-reading both the manifest and the
|
||||
//! scaler sidecar so feature-changing uploads land without a restart.
|
||||
//! Events inside the staging subdirectory are filtered so partial uploads
|
||||
//! can't flicker the UI through transient `Error` states.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use macros::log;
|
||||
use notify::{Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::sleep;
|
||||
|
||||
use super::adapter::ModelSourceState;
|
||||
use super::inference::Inference;
|
||||
use super::model_loader::MLModels;
|
||||
use super::model_loader::build_adapter;
|
||||
use crate::infrastructure::app_config::AppConfig;
|
||||
use crate::model::config::constants::{MANIFEST_FILENAME, MODELS_DIR, STAGING_SUBDIR};
|
||||
use crate::model::detection::model_source::ModelInfo;
|
||||
use crate::model::error::ml::MLError;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
|
||||
/// Debounce window: wait for both AE + classifier to land before reloading.
|
||||
/// Debounce window: wait for both manifest and ONNX to land before reloading.
|
||||
const DEBOUNCE_SECS: u64 = 5;
|
||||
|
||||
/// Watch the models/ directory for .onnx file changes and hot-reload into the inference pipeline.
|
||||
pub struct ModelWatcher {
|
||||
inference: Arc<Inference>,
|
||||
app_config: Arc<AppConfig>,
|
||||
inference_config: Arc<MLInferenceConfig>,
|
||||
}
|
||||
|
||||
impl ModelWatcher {
|
||||
pub fn new(
|
||||
inference: Arc<Inference>,
|
||||
app_config: Arc<AppConfig>,
|
||||
inference_config: Arc<MLInferenceConfig>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inference,
|
||||
app_config,
|
||||
inference_config,
|
||||
}
|
||||
pub fn new(inference: Arc<Inference>, app_config: Arc<AppConfig>) -> Self {
|
||||
Self { inference, app_config }
|
||||
}
|
||||
|
||||
/// Start watching in a background task.
|
||||
pub fn start(self) {
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = self.run().await {
|
||||
@ -50,7 +48,7 @@ impl ModelWatcher {
|
||||
}
|
||||
|
||||
async fn run(self) -> Result<(), MLError> {
|
||||
let models_dir = PathBuf::from("models");
|
||||
let models_dir = PathBuf::from(MODELS_DIR);
|
||||
if !models_dir.exists() {
|
||||
log!(MLLog::InferenceFailed(
|
||||
"ModelWatcher".to_string(),
|
||||
@ -60,23 +58,16 @@ impl ModelWatcher {
|
||||
}
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<()>(16);
|
||||
|
||||
// notify watcher runs on a blocking thread — forward events via channel
|
||||
let _watcher = Self::spawn_watcher(models_dir, tx)?;
|
||||
|
||||
log!(MLLog::ModelWatcherStarted);
|
||||
|
||||
loop {
|
||||
// Wait for first event
|
||||
if rx.recv().await.is_none() {
|
||||
break;
|
||||
}
|
||||
|
||||
// Debounce: drain any additional events within the window
|
||||
sleep(Duration::from_secs(DEBOUNCE_SECS)).await;
|
||||
while rx.try_recv().is_ok() {}
|
||||
|
||||
// Attempt reload
|
||||
self.try_reload();
|
||||
}
|
||||
|
||||
@ -86,37 +77,135 @@ impl ModelWatcher {
|
||||
fn spawn_watcher(models_dir: PathBuf, tx: mpsc::Sender<()>) -> Result<RecommendedWatcher, MLError> {
|
||||
let mut watcher = notify::recommended_watcher(move |res: Result<Event, notify::Error>| {
|
||||
if let Ok(event) = res {
|
||||
let dominated = matches!(event.kind, EventKind::Create(_) | EventKind::Modify(_));
|
||||
let has_onnx = event
|
||||
.paths
|
||||
.iter()
|
||||
.any(|p| p.extension().is_some_and(|ext| ext == "onnx"));
|
||||
if dominated && has_onnx {
|
||||
let _ = tx.blocking_send(());
|
||||
if !is_relevant_event(&event) {
|
||||
return;
|
||||
}
|
||||
let _ = tx.blocking_send(());
|
||||
}
|
||||
})
|
||||
.map_err(MLError::ModelWatcherFailed)?;
|
||||
|
||||
// Recursive watch so the staging filter gets exercised — otherwise a
|
||||
// drop-in to `.staging/` wouldn't trigger notify at all on some FSes.
|
||||
watcher
|
||||
.watch(&models_dir, RecursiveMode::NonRecursive)
|
||||
.watch(&models_dir, RecursiveMode::Recursive)
|
||||
.map_err(MLError::ModelWatcherFailed)?;
|
||||
Ok(watcher)
|
||||
}
|
||||
|
||||
/// Full reload: manifest presence check → config re-parse → adapter build →
|
||||
/// atomic state swap. Any failure lands the pipeline in `Error` rather
|
||||
/// than crashing.
|
||||
fn try_reload(&self) {
|
||||
let batch_size = self.app_config.inference.inference_batch_size;
|
||||
let manifest_path = PathBuf::from(MODELS_DIR).join(MANIFEST_FILENAME);
|
||||
|
||||
// Path 1 — manifest disappeared: transition to Dormant.
|
||||
if !manifest_path.exists() {
|
||||
log!(MLLog::ModelReloadStarting);
|
||||
self.inference.swap_state(ModelSourceState::Dormant);
|
||||
log!(MLLog::ModelReloadSuccess);
|
||||
return;
|
||||
}
|
||||
|
||||
// Path 2 — manifest present: re-read + rebuild adapter.
|
||||
log!(MLLog::ModelReloadStarting);
|
||||
|
||||
match MLModels::load_models(&self.app_config, &self.inference_config, batch_size) {
|
||||
Ok(new_models) => {
|
||||
self.inference.swap_models(Arc::new(new_models));
|
||||
log!(MLLog::ModelReloadSuccess);
|
||||
}
|
||||
let batch_size = self.app_config.inference.inference_batch_size;
|
||||
match MLInferenceConfig::from_manifest_with_sidecar(&manifest_path) {
|
||||
Ok((config, manifest)) => match build_adapter(&manifest, Some(&manifest_path), &config, batch_size) {
|
||||
Ok(adapter) => {
|
||||
let info = ModelInfo::new(
|
||||
manifest.name.clone(),
|
||||
manifest.adapter.as_str().to_string(),
|
||||
manifest.features.len(),
|
||||
);
|
||||
self.inference.swap_state(ModelSourceState::Active { adapter, info });
|
||||
log!(MLLog::ModelReloadSuccess);
|
||||
}
|
||||
Err(e) => {
|
||||
self.record_error(e.to_string(), Some(manifest_path.clone()));
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
log!(MLLog::ModelReloadFailed(e.to_string()));
|
||||
self.record_error(e.to_string(), Some(manifest_path.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn record_error(&self, msg: String, last_attempted_path: Option<PathBuf>) {
|
||||
log!(MLLog::ModelReloadFailed(msg.clone()));
|
||||
self.inference.swap_state(ModelSourceState::Error {
|
||||
msg,
|
||||
since: SystemTime::now(),
|
||||
last_attempted_path,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Inbound event filter. Ignore `.staging/` paths entirely; pass through
|
||||
/// `.onnx` / `.yaml` / `.yml` / `.json` changes in `models/`.
|
||||
fn is_relevant_event(event: &Event) -> bool {
|
||||
// Only Create / Modify events trigger a reload; renames / removes would
|
||||
// also surface but debouncing handles both equally well.
|
||||
if !matches!(event.kind, EventKind::Create(_) | EventKind::Modify(_)) {
|
||||
return false;
|
||||
}
|
||||
event.paths.iter().any(|p| {
|
||||
if path_is_inside_staging(p) {
|
||||
return false;
|
||||
}
|
||||
matches!(
|
||||
p.extension().and_then(|e| e.to_str()),
|
||||
Some("onnx") | Some("yaml") | Some("yml") | Some("json")
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn path_is_inside_staging(path: &Path) -> bool {
|
||||
path.components().any(|c| c.as_os_str() == STAGING_SUBDIR)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
use notify::event::{CreateKind, ModifyKind};
|
||||
|
||||
#[test]
|
||||
fn staging_paths_are_filtered() {
|
||||
assert!(path_is_inside_staging(&PathBuf::from("models/.staging/bad.onnx")));
|
||||
assert!(path_is_inside_staging(&PathBuf::from(
|
||||
"/tmp/models/.staging/sub/x.yaml"
|
||||
)));
|
||||
assert!(!path_is_inside_staging(&PathBuf::from("models/good.onnx")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_relevant_extensions_rejected() {
|
||||
let event = Event {
|
||||
kind: EventKind::Create(CreateKind::Any),
|
||||
paths: vec![PathBuf::from("models/readme.md")],
|
||||
attrs: Default::default(),
|
||||
};
|
||||
assert!(!is_relevant_event(&event));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn onnx_outside_staging_accepted() {
|
||||
let event = Event {
|
||||
kind: EventKind::Create(CreateKind::Any),
|
||||
paths: vec![PathBuf::from("models/foo.onnx")],
|
||||
attrs: Default::default(),
|
||||
};
|
||||
assert!(is_relevant_event(&event));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn staging_paths_always_rejected() {
|
||||
let event = Event {
|
||||
kind: EventKind::Modify(ModifyKind::Any),
|
||||
paths: vec![PathBuf::from("models/.staging/partial.onnx")],
|
||||
attrs: Default::default(),
|
||||
};
|
||||
assert!(!is_relevant_event(&event));
|
||||
}
|
||||
}
|
||||
|
||||
@ -20,6 +20,7 @@ use crate::interface::port::notification::AlertNotifier;
|
||||
use crate::interface::port::rate_limit_api::RateLimitPort;
|
||||
use crate::interface::port::secret_store::SecretStorePort;
|
||||
use crate::model::config::constants::MAX_PENDING_UNBLOCK_RETRIES;
|
||||
use crate::model::detection::attack_type::canonical_from_str;
|
||||
use crate::model::error::Error;
|
||||
use crate::model::error::soar::SoarError;
|
||||
use crate::model::event::ThreatDetectedEvent;
|
||||
@ -59,7 +60,9 @@ pub struct SoarEngine {
|
||||
pub(super) geoip: Option<Arc<GeoIpService>>,
|
||||
/// Optional rate limit config for adjust_rate_limit action.
|
||||
pub(super) rate_limit: Option<Arc<dyn RateLimitPort>>,
|
||||
/// Lock to serialize rate limit read-save-write sequences (Item 6: atomicity).
|
||||
/// Serializes rate-limit read-save-write sequences so concurrent SOAR
|
||||
/// actions can't race the `soar_rate_limit_*` DB settings into an
|
||||
/// inconsistent pair.
|
||||
pub(super) rate_limit_lock: TokioMutex<()>,
|
||||
/// Cached enforce level: Monitor=0, MlOnly=1, Enforce=2.
|
||||
pub(super) enforce_level_cache: Arc<AtomicU8>,
|
||||
@ -120,7 +123,7 @@ impl SoarEngine {
|
||||
// Check if this row belongs to the same playbook as the last one
|
||||
let needs_new = playbooks.last().is_none_or(|last| last.id != pb_id);
|
||||
if needs_new {
|
||||
let _ = threshold; // stored in DB but no longer used at runtime
|
||||
let _ = threshold; // persisted for schema stability; runtime gating comes from the Condition rows
|
||||
playbooks.push(Playbook {
|
||||
id: pb_id,
|
||||
name,
|
||||
@ -165,6 +168,9 @@ impl SoarEngine {
|
||||
}
|
||||
ConditionType::RepeatOffender => operator == "==",
|
||||
ConditionType::Frequency => operator == ">=",
|
||||
ConditionType::MultiSourceMin => operator == ">=",
|
||||
ConditionType::SingleSourceHigh => operator == ">=",
|
||||
ConditionType::FusedConfidenceAbove => matches!(operator.as_str(), ">=" | "<="),
|
||||
};
|
||||
if !valid {
|
||||
log!(SoarLog::InvalidConditionOperator(
|
||||
@ -183,6 +189,18 @@ impl SoarEngine {
|
||||
}
|
||||
}
|
||||
|
||||
// Warn on playbooks whose trigger_event isn't in the canonical
|
||||
// dictionary — those won't ever match a fused event and usually
|
||||
// signal a typo or stale pre-canonicalization playbook.
|
||||
for pb in &playbooks {
|
||||
if canonical_from_str(&pb.trigger_event).is_none() {
|
||||
log!(SoarLog::NonCanonicalTriggerEvent(
|
||||
pb.name.clone(),
|
||||
pb.trigger_event.clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
*self.playbooks.write() = playbooks;
|
||||
|
||||
// Load admin whitelist
|
||||
@ -334,7 +352,6 @@ impl SoarEngine {
|
||||
/// Restore original rate limits if the TTL has expired.
|
||||
/// Called by TTL scheduler on each sweep.
|
||||
pub async fn check_rate_limit_restoration(&self) -> Result<(), Error> {
|
||||
// Acquire lock to serialize rate limit read-save-write (Item 6: atomicity)
|
||||
let _guard = self.rate_limit_lock.lock().await;
|
||||
|
||||
let expires_str = match self
|
||||
@ -484,6 +501,8 @@ mod tests {
|
||||
geoip_country: None,
|
||||
is_repeat_offender: false,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
@ -519,6 +538,8 @@ mod tests {
|
||||
geoip_country: None,
|
||||
is_repeat_offender: false,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
@ -601,6 +622,8 @@ mod tests {
|
||||
geoip_country: None,
|
||||
is_repeat_offender: false,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
@ -636,6 +659,8 @@ mod tests {
|
||||
geoip_country: None,
|
||||
is_repeat_offender: false,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
@ -681,6 +706,8 @@ mod tests {
|
||||
geoip_country: None,
|
||||
is_repeat_offender: false,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
@ -757,6 +784,8 @@ mod tests {
|
||||
geoip_country: country.map(|s| s.to_string()),
|
||||
is_repeat_offender: repeat,
|
||||
sources: vec![DetectionSource::ML],
|
||||
active_source_count: 1,
|
||||
fused_confidence: 0.95,
|
||||
ae_score: 0.0,
|
||||
anomaly_score: 0.0,
|
||||
c2_score: 0.0,
|
||||
|
||||
@ -6,17 +6,29 @@
|
||||
//! but none of them touch anything outside those fields.
|
||||
|
||||
use std::net::IpAddr;
|
||||
use std::str::FromStr;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use macros::log;
|
||||
|
||||
use crate::core::soar::engine::SoarEngine;
|
||||
use crate::model::event::ThreatDetectedEvent;
|
||||
use crate::model::event::{DetectionSource, ThreatDetectedEvent};
|
||||
use crate::model::log::soar::SoarLog;
|
||||
use crate::model::soar::condition::{ConditionType, PlaybookCondition};
|
||||
use crate::model::soar::playbook::Playbook;
|
||||
|
||||
/// Default frequency-condition window when the playbook omits `value2`.
|
||||
const DEFAULT_FREQUENCY_WINDOW_SECS: u64 = 60;
|
||||
|
||||
/// Default minimum confidence for `SingleSourceHigh` when the playbook
|
||||
/// omits `value2`. Conservative enough that ad-hoc solo playbooks don't
|
||||
/// auto-block noisy single-source hits.
|
||||
const DEFAULT_SINGLE_SOURCE_HIGH_MIN_CONFIDENCE: f32 = 0.95;
|
||||
|
||||
/// Default expiry for cooldown cleanup when no playbook has a cooldown set.
|
||||
const DEFAULT_COOLDOWN_EXPIRY_SECS: u64 = 3600;
|
||||
|
||||
impl SoarEngine {
|
||||
/// Find playbooks matching the event via trigger_event + multi-condition AND logic.
|
||||
pub(super) fn find_matching_playbooks(&self, event: &ThreatDetectedEvent) -> Vec<Playbook> {
|
||||
@ -142,7 +154,11 @@ impl SoarEngine {
|
||||
Ok(v) => v,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let window_secs = cond.value2.as_ref().and_then(|s| s.parse::<u64>().ok()).unwrap_or(60);
|
||||
let window_secs = cond
|
||||
.value2
|
||||
.as_ref()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(DEFAULT_FREQUENCY_WINDOW_SECS);
|
||||
let count = self
|
||||
.frequency_tracker
|
||||
.record_and_count(pb.id, &event.source_ip, window_secs);
|
||||
@ -152,6 +168,70 @@ impl SoarEngine {
|
||||
}
|
||||
met
|
||||
}
|
||||
ConditionType::MultiSourceMin => {
|
||||
let required = match cond.value.parse::<usize>() {
|
||||
Ok(v) => v,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let met = event.active_source_count >= required;
|
||||
if !met {
|
||||
log!(SoarLog::ConditionNotMet(
|
||||
"multi_source_min".to_string(),
|
||||
pb.name.clone(),
|
||||
format!("{} / {}", event.active_source_count, required),
|
||||
));
|
||||
}
|
||||
met
|
||||
}
|
||||
ConditionType::SingleSourceHigh => {
|
||||
// `value` names the target source (canonical Display form
|
||||
// or a common alias); `value2` is the minimum confidence.
|
||||
let target_source = match DetectionSource::from_str(&cond.value) {
|
||||
Ok(s) => s,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let min_conf = cond
|
||||
.value2
|
||||
.as_ref()
|
||||
.and_then(|s| s.parse::<f32>().ok())
|
||||
.unwrap_or(DEFAULT_SINGLE_SOURCE_HIGH_MIN_CONFIDENCE);
|
||||
// Solo = exactly one contributing source AND it matches the
|
||||
// target source AND confidence clears the escape-hatch bar.
|
||||
let solo_match =
|
||||
event.active_source_count == 1 && event.sources.len() == 1 && event.sources[0] == target_source;
|
||||
let conf_met = event.confidence >= min_conf;
|
||||
let met = solo_match && conf_met;
|
||||
if !met {
|
||||
log!(SoarLog::ConditionNotMet(
|
||||
"single_source_high".to_string(),
|
||||
pb.name.clone(),
|
||||
format!(
|
||||
"sources={:?} count={} conf={:.3} target={} need_conf>={:.3}",
|
||||
event.sources, event.active_source_count, event.confidence, cond.value, min_conf
|
||||
),
|
||||
));
|
||||
}
|
||||
met
|
||||
}
|
||||
ConditionType::FusedConfidenceAbove => {
|
||||
let threshold = match cond.value.parse::<f32>() {
|
||||
Ok(v) => v,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let met = if cond.operator == "<=" {
|
||||
event.fused_confidence <= threshold
|
||||
} else {
|
||||
event.fused_confidence >= threshold
|
||||
};
|
||||
if !met {
|
||||
log!(SoarLog::ConditionNotMet(
|
||||
"fused_confidence_above".to_string(),
|
||||
pb.name.clone(),
|
||||
format!("{:.3} vs {:.3}", event.fused_confidence, threshold),
|
||||
));
|
||||
}
|
||||
met
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -178,9 +258,13 @@ impl SoarEngine {
|
||||
pub fn cleanup_expired_cooldowns(&self) {
|
||||
let max_cooldown_secs = {
|
||||
let playbooks = self.playbooks.read();
|
||||
playbooks.iter().map(|p| p.cooldown_secs as u64).max().unwrap_or(3600)
|
||||
playbooks
|
||||
.iter()
|
||||
.map(|p| p.cooldown_secs as u64)
|
||||
.max()
|
||||
.unwrap_or(DEFAULT_COOLDOWN_EXPIRY_SECS)
|
||||
};
|
||||
let expiry = Duration::from_secs(max_cooldown_secs.saturating_mul(2).max(3600));
|
||||
let expiry = Duration::from_secs(max_cooldown_secs.saturating_mul(2).max(DEFAULT_COOLDOWN_EXPIRY_SECS));
|
||||
let before = self.cooldowns.len();
|
||||
self.cooldowns.retain(|_, instant| instant.elapsed() < expiry);
|
||||
let removed = before.saturating_sub(self.cooldowns.len());
|
||||
|
||||
@ -1,24 +1,30 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use crossbeam::queue::SegQueue;
|
||||
use macros::log;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use crate::core::ml::adapter::ModelSourceState;
|
||||
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::inference::Inference;
|
||||
use crate::core::ml::manifest::ModelManifest;
|
||||
use crate::core::ml::model_loader::build_adapter;
|
||||
use crate::core::ml::traffic_logger::TrafficLogger;
|
||||
use crate::infrastructure::app_config::AppConfig;
|
||||
use crate::infrastructure::health::SystemHealth;
|
||||
use crate::infrastructure::statistics::FlowStatistics;
|
||||
use crate::model::config::constants::{MANIFEST_FILENAME, MODELS_DIR};
|
||||
use crate::model::detection::flow_features::FlowFeatures;
|
||||
use crate::model::detection::ml_detection::EngineConfig;
|
||||
use crate::model::detection::model_source::ModelInfo;
|
||||
use crate::model::error::Error;
|
||||
use crate::model::error::misc::MiscError;
|
||||
use crate::model::error::system::SystemError;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::log::system::SystemLog;
|
||||
use crate::model::system::config::MLInferenceConfig;
|
||||
use crate::model::system::health::EbpfHealth;
|
||||
@ -29,7 +35,7 @@ use crate::model::system::health::EbpfHealth;
|
||||
pub struct AppServices {
|
||||
pub health: Arc<SystemHealth>,
|
||||
pub ml_alert: Arc<MLAlert>,
|
||||
pub ml_models: Arc<MLModels>,
|
||||
pub ml_inference: Arc<Inference>,
|
||||
pub ml_engine: Arc<Engine>,
|
||||
pub flow_statistics: Arc<FlowStatistics>,
|
||||
shutdowns: SegQueue<oneshot::Sender<()>>,
|
||||
@ -46,25 +52,45 @@ impl AppServices {
|
||||
let health = SystemHealth::new(app_config.clone(), ebpf_health)?;
|
||||
|
||||
let batch_size = app_config.inference.inference_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)?
|
||||
|
||||
let initial_state = match ml_manifest.as_ref() {
|
||||
Some(manifest) => {
|
||||
let manifest_path = PathBuf::from(MODELS_DIR).join(MANIFEST_FILENAME);
|
||||
match build_adapter(manifest, Some(&manifest_path), &inference_config, batch_size) {
|
||||
Ok(adapter) => {
|
||||
let info = ModelInfo::new(
|
||||
manifest.name.clone(),
|
||||
manifest.adapter.as_str().to_string(),
|
||||
manifest.features.len(),
|
||||
);
|
||||
log!(MLLog::ModelsLoaded(format!(
|
||||
"{} ({}) — {} features, {} labels",
|
||||
manifest.name,
|
||||
manifest.adapter.as_str(),
|
||||
manifest.features.len(),
|
||||
manifest.labels.len()
|
||||
)));
|
||||
ModelSourceState::Active { adapter, info }
|
||||
}
|
||||
Err(e) => {
|
||||
log!(MLLog::ModelReloadFailed(e.to_string()));
|
||||
ModelSourceState::Error {
|
||||
msg: e.to_string(),
|
||||
since: SystemTime::now(),
|
||||
last_attempted_path: Some(manifest_path),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => MLModels::load_models(&app_config, &inference_config, batch_size)?,
|
||||
});
|
||||
None => {
|
||||
log!(MLLog::ModelsLoaded(
|
||||
"no manifest present — ML source dormant".to_string()
|
||||
));
|
||||
ModelSourceState::Dormant
|
||||
}
|
||||
};
|
||||
|
||||
let ml_inference = Arc::new(Inference::new(initial_state, inference_config.clone()));
|
||||
let ml_alert = Arc::new(MLAlert::new());
|
||||
|
||||
let traffic_logger = if app_config.inference.traffic_logging_mode {
|
||||
@ -88,8 +114,7 @@ impl AppServices {
|
||||
};
|
||||
|
||||
let ml_engine = Arc::new(Engine::new(
|
||||
ml_models.clone(),
|
||||
inference_config.clone(),
|
||||
ml_inference.clone(),
|
||||
ml_alert.clone(),
|
||||
drift_detector,
|
||||
engine_config,
|
||||
@ -102,7 +127,7 @@ impl AppServices {
|
||||
Ok(Self {
|
||||
health: Arc::new(health),
|
||||
ml_alert,
|
||||
ml_models,
|
||||
ml_inference,
|
||||
ml_engine,
|
||||
flow_statistics,
|
||||
shutdowns: SegQueue::new(),
|
||||
|
||||
@ -156,12 +156,15 @@ impl System {
|
||||
pub async fn run(&mut self) -> Result<ShutdownMode, Error> {
|
||||
log!(SystemLog::Initializing);
|
||||
|
||||
log!(MLLog::ModelsLoaded(
|
||||
self.app_services.ml_models.get_model_info("deep_autoencoder")
|
||||
));
|
||||
log!(MLLog::ModelsLoaded(
|
||||
self.app_services.ml_models.get_model_info("classifier")
|
||||
));
|
||||
// ML source state snapshot. Day 1 with no manifest renders as Dormant
|
||||
// — the rest of the stack still runs (3-source fusion).
|
||||
{
|
||||
let status = self.app_services.ml_inference.current_status();
|
||||
match serde_json::to_string(&status) {
|
||||
Ok(s) => log!(MLLog::ModelsLoaded(s)),
|
||||
Err(e) => log!(MLLog::ModelsLoaded(format!("<unserializable status: {e}>"))),
|
||||
}
|
||||
}
|
||||
log!(MLLog::ConfigLoaded(
|
||||
self.inference_config.num_ae_features(),
|
||||
self.inference_config.num_attack_types(),
|
||||
@ -267,11 +270,7 @@ impl System {
|
||||
});
|
||||
|
||||
// Start model hot-reload watcher (monitors models/ for .onnx changes)
|
||||
let model_watcher = ModelWatcher::new(
|
||||
self.app_services.ml_engine.inference_pipeline().clone(),
|
||||
self.app_config.clone(),
|
||||
self.inference_config.clone(),
|
||||
);
|
||||
let model_watcher = ModelWatcher::new(self.app_services.ml_inference.clone(), self.app_config.clone());
|
||||
model_watcher.start();
|
||||
|
||||
// Initialize force_https flag from DB setting
|
||||
|
||||
@ -14,6 +14,11 @@ pub const FLOW_BULK_MIN_BYTES: u64 = 1000;
|
||||
pub const FLOW_IDLE_TIMEOUT_US: u64 = 120_000_000;
|
||||
pub const FLOW_TERMINATED_TIMEOUT_US: u64 = 5_000_000;
|
||||
|
||||
// ── ML Model Directory ─────────────────────────────────────────────
|
||||
pub const MODELS_DIR: &str = "models";
|
||||
pub const MANIFEST_FILENAME: &str = "manifest.yaml";
|
||||
pub const STAGING_SUBDIR: &str = ".staging";
|
||||
|
||||
// ── Notification ───────────────────────────────────────────────────
|
||||
pub const TELEGRAM_MAX_RETRIES: u32 = 2;
|
||||
|
||||
|
||||
327
net-guardia/src/model/detection/attack_type.rs
Normal file
327
net-guardia/src/model/detection/attack_type.rs
Normal file
@ -0,0 +1,327 @@
|
||||
//! Canonical attack-type dictionary and cross-source translator.
|
||||
//!
|
||||
//! Fusion v1 requires that Suricata / ML / CV (Beaconing) / Graph (Correlation)
|
||||
//! use a shared vocabulary — otherwise the orchestrator's dedup key
|
||||
//! `(source_ip, attack_type)` never collides across sources, and the
|
||||
//! cross-source "sources agreed" fusion signal is impossible.
|
||||
//!
|
||||
//! The 13 canonical types below are the v1 seed set. Each source ships a
|
||||
//! translation table from its own raw labels (Suricata classtype, ML class
|
||||
//! name, Beaconing tag, Correlation sub-type) into the canonical vocabulary.
|
||||
//! Unknown labels land in `Unknown` — a valid dedup bucket that still
|
||||
//! participates in fusion.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::model::event::DetectionSource;
|
||||
|
||||
/// The v1 seed dictionary — 13 canonical attack types every detection source
|
||||
/// maps into. New types may be added without breaking change; renaming or
|
||||
/// removing one IS a breaking change (dedup keys drift).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CanonicalAttackType {
|
||||
BruteForce,
|
||||
PortScan,
|
||||
C2Beacon,
|
||||
DnsTunnel,
|
||||
SqlInjection,
|
||||
Xss,
|
||||
Exploit,
|
||||
LateralMovement,
|
||||
Reconnaissance,
|
||||
Cryptomining,
|
||||
DosDdos,
|
||||
BotActivity,
|
||||
/// Fallback bucket. Still a valid dedup key — two sources firing
|
||||
/// `Unknown` on the same src_ip within the fusion window DO fuse.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl CanonicalAttackType {
|
||||
/// Stable wire-format string. Must match `DedupKey.attack_type` exactly
|
||||
/// across releases — renaming breaks dedup on in-flight alerts.
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::BruteForce => "brute_force",
|
||||
Self::PortScan => "port_scan",
|
||||
Self::C2Beacon => "c2_beacon",
|
||||
Self::DnsTunnel => "dns_tunnel",
|
||||
Self::SqlInjection => "sql_injection",
|
||||
Self::Xss => "xss",
|
||||
Self::Exploit => "exploit",
|
||||
Self::LateralMovement => "lateral_movement",
|
||||
Self::Reconnaissance => "reconnaissance",
|
||||
Self::Cryptomining => "cryptomining",
|
||||
Self::DosDdos => "dos_ddos",
|
||||
Self::BotActivity => "bot_activity",
|
||||
Self::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
/// All 13 canonical values, in declaration order. Used by BYO-contract
|
||||
/// docs + CI consistency check.
|
||||
pub const ALL: &'static [Self] = &[
|
||||
Self::BruteForce,
|
||||
Self::PortScan,
|
||||
Self::C2Beacon,
|
||||
Self::DnsTunnel,
|
||||
Self::SqlInjection,
|
||||
Self::Xss,
|
||||
Self::Exploit,
|
||||
Self::LateralMovement,
|
||||
Self::Reconnaissance,
|
||||
Self::Cryptomining,
|
||||
Self::DosDdos,
|
||||
Self::BotActivity,
|
||||
Self::Unknown,
|
||||
];
|
||||
}
|
||||
|
||||
impl fmt::Display for CanonicalAttackType {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
/// Reverse lookup: canonical wire string → enum. Returns `None` when the
|
||||
/// input isn't one of the 13 seeds (caller typically treats this as an
|
||||
/// upstream bug, not an `Unknown` bucket).
|
||||
pub fn canonical_from_str(s: &str) -> Option<CanonicalAttackType> {
|
||||
CanonicalAttackType::ALL.iter().copied().find(|c| c.as_str() == s)
|
||||
}
|
||||
|
||||
/// Translate a source-specific raw label into the canonical dictionary.
|
||||
///
|
||||
/// Never fails — unknown labels map to `CanonicalAttackType::Unknown` so
|
||||
/// the event still lands in a valid dedup bucket. The `raw_label`
|
||||
/// comparison is case-insensitive and trims whitespace; callers that pass
|
||||
/// user-visible labels verbatim don't need to pre-normalize.
|
||||
pub fn translate(source: DetectionSource, raw_label: &str) -> CanonicalAttackType {
|
||||
let normalized = raw_label.trim().to_ascii_lowercase();
|
||||
match source {
|
||||
DetectionSource::ML => translate_ml(&normalized),
|
||||
DetectionSource::Beaconing => translate_beaconing(&normalized),
|
||||
DetectionSource::Correlation => translate_correlation(&normalized),
|
||||
DetectionSource::Suricata => translate_suricata(&normalized),
|
||||
}
|
||||
}
|
||||
|
||||
/// Map ML classifier class names into the canonical dictionary. Case-insensitive;
|
||||
/// covers the baseline class vocabulary.
|
||||
fn translate_ml(label: &str) -> CanonicalAttackType {
|
||||
match label {
|
||||
"brute force" | "brute_force" | "bruteforce" => CanonicalAttackType::BruteForce,
|
||||
"reconnaissance" | "recon" | "port_scan" | "portscan" => CanonicalAttackType::Reconnaissance,
|
||||
"c2 communication" | "c2" | "c2_beacon" | "command_and_control" => CanonicalAttackType::C2Beacon,
|
||||
"dns tunneling" | "dns_tunnel" | "dns tunnel" => CanonicalAttackType::DnsTunnel,
|
||||
"sql injection" | "sql_injection" | "sqli" => CanonicalAttackType::SqlInjection,
|
||||
"xss" | "cross_site_scripting" => CanonicalAttackType::Xss,
|
||||
"web attack" | "web_attack" => CanonicalAttackType::Exploit,
|
||||
"exploitation" | "exploit" => CanonicalAttackType::Exploit,
|
||||
"dos/ddos" | "dos_ddos" | "ddos" | "dos" => CanonicalAttackType::DosDdos,
|
||||
"cryptomining" | "cryptocurrency_mining" | "mining" => CanonicalAttackType::Cryptomining,
|
||||
"bot" | "bot_activity" | "botnet" | "malware" => CanonicalAttackType::BotActivity,
|
||||
"lateral movement" | "lateral_movement" => CanonicalAttackType::LateralMovement,
|
||||
// "Normal" is not a threat — translators should not see it, but if they do,
|
||||
// fall through to Unknown rather than panicking.
|
||||
_ => CanonicalAttackType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
/// Beaconing (Layer 2 CV) only produces C2-style temporal beacons in v1.
|
||||
/// Sub-tags (e.g. "c2_beacon", "heartbeat") all collapse here.
|
||||
fn translate_beaconing(_label: &str) -> CanonicalAttackType {
|
||||
CanonicalAttackType::C2Beacon
|
||||
}
|
||||
|
||||
/// Correlation (Layer 3 graph) splits across scan / lateral / botnet.
|
||||
fn translate_correlation(label: &str) -> CanonicalAttackType {
|
||||
match label {
|
||||
"scan" | "port_scan" | "reconnaissance" => CanonicalAttackType::PortScan,
|
||||
"lateral" | "lateral_movement" => CanonicalAttackType::LateralMovement,
|
||||
"botnet" | "bot" | "bot_activity" => CanonicalAttackType::BotActivity,
|
||||
_ => CanonicalAttackType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
/// Map Suricata classtypes (`eve.json.alert.category`, not sid) into the
|
||||
/// canonical dictionary. We match on classtype because sid numbering isn't
|
||||
/// stable across rule packs; classtype is part of the rule DSL and stable
|
||||
/// across ET Open / Talos releases.
|
||||
fn translate_suricata(label: &str) -> CanonicalAttackType {
|
||||
// classtype strings come lowercase+trimmed from `translate`
|
||||
match label {
|
||||
// Scan / reconnaissance
|
||||
"attempted-recon" | "network-scan" | "misc-activity" => CanonicalAttackType::Reconnaissance,
|
||||
// Exploits / admin compromise
|
||||
"attempted-admin" | "successful-admin" | "attempted-user" | "successful-user" | "shellcode-detect"
|
||||
| "attempted-exploit" => CanonicalAttackType::Exploit,
|
||||
// Web-application attacks
|
||||
"web-application-attack" => CanonicalAttackType::Exploit,
|
||||
"web-application-activity" => CanonicalAttackType::Exploit,
|
||||
// SQL injection is usually emitted as web-application-attack, but some
|
||||
// rule packs use "sql-injection" directly.
|
||||
"sql-injection" => CanonicalAttackType::SqlInjection,
|
||||
// XSS — same note as SQLi
|
||||
"xss" | "cross-site-scripting" => CanonicalAttackType::Xss,
|
||||
// DoS / DDoS
|
||||
"attempted-dos" | "successful-dos" | "denial-of-service" => CanonicalAttackType::DosDdos,
|
||||
// Trojan / malware / C2
|
||||
"trojan-activity" | "malware-cnc" | "command-and-control" => CanonicalAttackType::C2Beacon,
|
||||
// Credential attacks
|
||||
"suspicious-login" | "unsuccessful-user" | "brute-force" => CanonicalAttackType::BruteForce,
|
||||
// Policy / Crypto miner
|
||||
"coin-mining" | "policy-violation" => CanonicalAttackType::Cryptomining,
|
||||
// DNS tunneling detections emitted by some ET Open rules
|
||||
"dns-tunnel" | "protocol-command-decode" => CanonicalAttackType::DnsTunnel,
|
||||
_ => CanonicalAttackType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn wire_format_roundtrip() {
|
||||
for canonical in CanonicalAttackType::ALL {
|
||||
let wire = canonical.as_str();
|
||||
let round = canonical_from_str(wire).expect("wire format must round-trip");
|
||||
assert_eq!(*canonical, round, "roundtrip mismatch for {wire}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_wire_string_returns_none() {
|
||||
assert!(canonical_from_str("this_attack_does_not_exist").is_none());
|
||||
assert!(canonical_from_str("").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ml_v10_class_names_translate() {
|
||||
// v10 ships 10 classes — every non-"Normal" class must hit a canonical entry.
|
||||
// "Normal" is a legitimate benign class and IS expected to return Unknown
|
||||
// because translators should not be called on benign flows in the first place.
|
||||
let cases = [
|
||||
("Brute Force", CanonicalAttackType::BruteForce),
|
||||
("Reconnaissance", CanonicalAttackType::Reconnaissance),
|
||||
("C2 Communication", CanonicalAttackType::C2Beacon),
|
||||
("DoS/DDoS", CanonicalAttackType::DosDdos),
|
||||
("Exploitation", CanonicalAttackType::Exploit),
|
||||
("Web Attack", CanonicalAttackType::Exploit),
|
||||
("Bot", CanonicalAttackType::BotActivity),
|
||||
("Malware", CanonicalAttackType::BotActivity),
|
||||
("Cryptomining", CanonicalAttackType::Cryptomining),
|
||||
];
|
||||
for (raw, expected) in cases {
|
||||
assert_eq!(
|
||||
translate(DetectionSource::ML, raw),
|
||||
expected,
|
||||
"ML class {raw:?} must translate to {expected:?}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn beaconing_always_c2() {
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Beaconing, "anything"),
|
||||
CanonicalAttackType::C2Beacon
|
||||
);
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Beaconing, "c2_beacon"),
|
||||
CanonicalAttackType::C2Beacon
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correlation_subtypes_split() {
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Correlation, "scan"),
|
||||
CanonicalAttackType::PortScan
|
||||
);
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Correlation, "lateral"),
|
||||
CanonicalAttackType::LateralMovement
|
||||
);
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Correlation, "botnet"),
|
||||
CanonicalAttackType::BotActivity
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suricata_classtype_mapping() {
|
||||
let cases = [
|
||||
("attempted-admin", CanonicalAttackType::Exploit),
|
||||
("web-application-attack", CanonicalAttackType::Exploit),
|
||||
("trojan-activity", CanonicalAttackType::C2Beacon),
|
||||
("attempted-recon", CanonicalAttackType::Reconnaissance),
|
||||
("attempted-dos", CanonicalAttackType::DosDdos),
|
||||
("coin-mining", CanonicalAttackType::Cryptomining),
|
||||
("brute-force", CanonicalAttackType::BruteForce),
|
||||
("sql-injection", CanonicalAttackType::SqlInjection),
|
||||
];
|
||||
for (raw, expected) in cases {
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Suricata, raw),
|
||||
expected,
|
||||
"Suricata classtype {raw:?} must translate to {expected:?}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_label_lands_in_unknown_bucket() {
|
||||
// Unknown is the fallback — MUST NOT panic, MUST be dedup-safe.
|
||||
for source in [
|
||||
DetectionSource::ML,
|
||||
DetectionSource::Correlation,
|
||||
DetectionSource::Suricata,
|
||||
] {
|
||||
assert_eq!(
|
||||
translate(source, "this_is_not_a_real_label"),
|
||||
CanonicalAttackType::Unknown,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn case_insensitive_and_whitespace_tolerant() {
|
||||
assert_eq!(
|
||||
translate(DetectionSource::ML, " BRUTE FORCE "),
|
||||
CanonicalAttackType::BruteForce,
|
||||
);
|
||||
assert_eq!(
|
||||
translate(DetectionSource::Suricata, "Attempted-Admin"),
|
||||
CanonicalAttackType::Exploit,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedup_key_non_collision_across_sources() {
|
||||
// The central invariant: two sources hitting the SAME attack on the
|
||||
// SAME src_ip produce the SAME canonical wire string, so dedup collides.
|
||||
let ml = translate(DetectionSource::ML, "Brute Force");
|
||||
let suricata = translate(DetectionSource::Suricata, "brute-force");
|
||||
assert_eq!(
|
||||
ml.as_str(),
|
||||
suricata.as_str(),
|
||||
"cross-source dedup key MUST match for the same canonical attack",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_canonical_have_unique_wire_strings() {
|
||||
use std::collections::HashSet;
|
||||
let strings: HashSet<&str> = CanonicalAttackType::ALL.iter().map(|c| c.as_str()).collect();
|
||||
assert_eq!(
|
||||
strings.len(),
|
||||
CanonicalAttackType::ALL.len(),
|
||||
"duplicate wire string in CanonicalAttackType::ALL",
|
||||
);
|
||||
}
|
||||
}
|
||||
@ -1,3 +1,5 @@
|
||||
pub mod attack_type;
|
||||
pub mod drift;
|
||||
pub mod flow_features;
|
||||
pub mod ml_detection;
|
||||
pub mod model_source;
|
||||
|
||||
113
net-guardia/src/model/detection/model_source.rs
Normal file
113
net-guardia/src/model/detection/model_source.rs
Normal file
@ -0,0 +1,113 @@
|
||||
//! Wire-level types for broadcasting ML source state to the frontend.
|
||||
//!
|
||||
//! Two layers:
|
||||
//! - `ModelInfo` — stable facts about a loaded model (name, adapter kind, feature count).
|
||||
//! - `ModelSourceStatus` — current state of the ML source: Dormant / Active / Error.
|
||||
//!
|
||||
//! The internal `core::ml::state::ModelSourceState` holds the actual model
|
||||
//! adapter plus this metadata; it converts into `ModelSourceStatus` for
|
||||
//! WebSocket / HTTP responses via `From`.
|
||||
|
||||
use std::time::SystemTime;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Public facts about a currently-loaded model. Drives the UI's ML Status panel.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelInfo {
|
||||
/// Human-readable manifest name (e.g., "netguardia-v10").
|
||||
pub name: String,
|
||||
/// One of "autoencoder_only" | "classifier_only" | "multi_task".
|
||||
pub adapter_kind: String,
|
||||
/// When this model became Active (unix epoch seconds).
|
||||
pub loaded_at_secs: u64,
|
||||
/// Number of features the manifest declares. Useful for UI "31 features" display.
|
||||
pub features_count: usize,
|
||||
/// Recent inference QPS (rolling window). Zero before first tick.
|
||||
pub qps_recent: f32,
|
||||
}
|
||||
|
||||
impl ModelInfo {
|
||||
pub fn new(name: String, adapter_kind: String, features_count: usize) -> Self {
|
||||
let loaded_at_secs = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
Self {
|
||||
name,
|
||||
adapter_kind,
|
||||
loaded_at_secs,
|
||||
features_count,
|
||||
qps_recent: 0.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wire-format ML source status. Broadcast to frontend; returned by
|
||||
/// `GET /api/models/current`. Serde-tagged so the frontend can discriminate
|
||||
/// on the `state` field.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "state", rename_all = "snake_case")]
|
||||
pub enum ModelSourceStatus {
|
||||
/// No model loaded. Day 1 default; fusion runs with 3 sources.
|
||||
Dormant,
|
||||
/// Model loaded and serving inference. `info` populates the UI card.
|
||||
Active { info: ModelInfo },
|
||||
/// Last load attempt failed. UI shows the reason in red.
|
||||
/// `since_secs` is unix epoch seconds; `last_attempted_path` is the
|
||||
/// file that failed (if any).
|
||||
Error {
|
||||
msg: String,
|
||||
since_secs: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
last_attempted_path: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ModelSourceStatus {
|
||||
pub fn is_active(&self) -> bool {
|
||||
matches!(self, Self::Active { .. })
|
||||
}
|
||||
|
||||
pub fn is_dormant(&self) -> bool {
|
||||
matches!(self, Self::Dormant)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_dormant() {
|
||||
let json = serde_json::to_string(&ModelSourceStatus::Dormant).unwrap();
|
||||
assert!(json.contains("\"state\":\"dormant\""));
|
||||
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
|
||||
assert!(back.is_dormant());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_active() {
|
||||
let status = ModelSourceStatus::Active {
|
||||
info: ModelInfo::new("netguardia-v10".into(), "multi_task".into(), 31),
|
||||
};
|
||||
let json = serde_json::to_string(&status).unwrap();
|
||||
assert!(json.contains("\"state\":\"active\""));
|
||||
assert!(json.contains("\"adapter_kind\":\"multi_task\""));
|
||||
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
|
||||
assert!(back.is_active());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_roundtrip_error() {
|
||||
let status = ModelSourceStatus::Error {
|
||||
msg: "feature mismatch".into(),
|
||||
since_secs: 1_700_000_000,
|
||||
last_attempted_path: Some("models/.staging/bad.onnx".into()),
|
||||
};
|
||||
let json = serde_json::to_string(&status).unwrap();
|
||||
assert!(json.contains("\"state\":\"error\""));
|
||||
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(back, ModelSourceStatus::Error { .. }));
|
||||
}
|
||||
}
|
||||
@ -1,12 +1,13 @@
|
||||
use std::fmt;
|
||||
use std::str::FromStr;
|
||||
|
||||
use crate::interface::communication::event::Event;
|
||||
|
||||
// -- Detection Source ---------------------------------------------------------
|
||||
|
||||
/// Identifies which detection subsystem produced a detection.
|
||||
/// Used for attribution tracking and future cross-source deduplication.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
/// Used for attribution tracking and cross-source deduplication.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum DetectionSource {
|
||||
ML,
|
||||
Correlation,
|
||||
@ -25,6 +26,22 @@ impl fmt::Display for DetectionSource {
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for DetectionSource {
|
||||
type Err = ();
|
||||
|
||||
/// Accepts the canonical `Display` form plus common aliases so playbook
|
||||
/// authors can write `"CV"` for Beaconing or `"Graph"` for Correlation.
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"ML" | "ml" => Ok(Self::ML),
|
||||
"Suricata" | "suricata" => Ok(Self::Suricata),
|
||||
"Beaconing" | "beaconing" | "CV" | "cv" => Ok(Self::Beaconing),
|
||||
"Correlation" | "correlation" | "Graph" | "graph" => Ok(Self::Correlation),
|
||||
_ => Err(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- Detection Event (internal pipeline) --------------------------------------
|
||||
|
||||
/// Raw detection from any source. Sent via mpsc channel to DetectionOrchestrator.
|
||||
@ -68,9 +85,19 @@ pub struct ThreatDetectedEvent {
|
||||
pub geoip_country: Option<String>,
|
||||
/// Whether this src_ip had a block action in the past 24h
|
||||
pub is_repeat_offender: bool,
|
||||
/// Which detection sources contributed to this threat (for attribution)
|
||||
/// Which detection sources contributed to this threat (for attribution).
|
||||
/// Single-source events have length 1; fused events have 2..=4.
|
||||
pub sources: Vec<DetectionSource>,
|
||||
/// Per-model scores for debugging false positives
|
||||
/// Number of distinct sources that contributed to this event.
|
||||
/// Used by SOAR `MultiSourceMin` / `SingleSourceHigh` conditions —
|
||||
/// counts unique sources, not per-source fires within the window.
|
||||
pub active_source_count: usize,
|
||||
/// Cross-source fused confidence (1 − ∏(1 − c_i)). Equal to
|
||||
/// `confidence` after the fusion engine runs; retained as a separate
|
||||
/// field so SOAR policies can discriminate "single source" from
|
||||
/// "fused multi-source" numerically identical confidences.
|
||||
pub fused_confidence: f32,
|
||||
/// Per-model scores for debugging false positives.
|
||||
pub ae_score: f32,
|
||||
pub anomaly_score: f32,
|
||||
pub c2_score: f32,
|
||||
|
||||
@ -44,5 +44,17 @@ loggable! {
|
||||
|
||||
#[error("Correlation cleanup: removed {removed} expired entries")]
|
||||
CorrelationCleanup { removed: usize } => tracing::Level::DEBUG,
|
||||
|
||||
#[error("Fusion: emitted {source_ip} {attack_type} fused={fused:.3} sources={count}")]
|
||||
FusionEmitted { source_ip: String, attack_type: String, fused: f32, count: usize } => tracing::Level::DEBUG,
|
||||
|
||||
#[error("Fusion: window evicted under LRU pressure ({key_src} {key_type})")]
|
||||
FusionWindowEvicted { key_src: String, key_type: String } => tracing::Level::WARN,
|
||||
|
||||
#[error("Fusion: failed to publish ThreatDetectedEvent: {err}")]
|
||||
FusionPublishFailed { err: String } => tracing::Level::ERROR,
|
||||
|
||||
#[error("Fusion: failed to publish AuditEvent: {err}")]
|
||||
FusionAuditPublishFailed { err: String } => tracing::Level::ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
@ -92,5 +92,8 @@ loggable! {
|
||||
|
||||
#[error("Invalid operator '{operator}' for condition type '{condition_type}' on playbook '{name}', condition skipped")]
|
||||
InvalidConditionOperator { name: String, condition_type: String, operator: String } => tracing::Level::WARN,
|
||||
|
||||
#[error("Playbook '{name}' uses non-canonical trigger_event '{trigger_event}' — cross-source fusion dedup may silently miss this rule")]
|
||||
NonCanonicalTriggerEvent { name: String, trigger_event: String } => tracing::Level::WARN,
|
||||
}
|
||||
}
|
||||
|
||||
@ -16,6 +16,17 @@ pub enum ConditionType {
|
||||
RepeatOffender,
|
||||
/// Frequency: N events from same source_ip within window_secs
|
||||
Frequency,
|
||||
/// Multi-source agreement: `event.active_source_count >= value`.
|
||||
MultiSourceMin,
|
||||
/// Solo high-confidence escape hatch: exactly one contributing source,
|
||||
/// and that source matches a specific name with confidence ≥ `value2`.
|
||||
/// Lets a single high-confidence signature-class detection block
|
||||
/// without waiting on peer agreement.
|
||||
SingleSourceHigh,
|
||||
/// Fused confidence above threshold — reads `event.fused_confidence`
|
||||
/// rather than the per-event `confidence`, so a single high-confidence
|
||||
/// event doesn't pass a threshold intended for multi-source agreement.
|
||||
FusedConfidenceAbove,
|
||||
}
|
||||
|
||||
impl fmt::Display for ConditionType {
|
||||
@ -26,6 +37,9 @@ impl fmt::Display for ConditionType {
|
||||
Self::IpPattern => write!(f, "ip_pattern"),
|
||||
Self::RepeatOffender => write!(f, "repeat_offender"),
|
||||
Self::Frequency => write!(f, "frequency"),
|
||||
Self::MultiSourceMin => write!(f, "multi_source_min"),
|
||||
Self::SingleSourceHigh => write!(f, "single_source_high"),
|
||||
Self::FusedConfidenceAbove => write!(f, "fused_confidence_above"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -40,6 +54,9 @@ impl FromStr for ConditionType {
|
||||
"ip_pattern" => Ok(Self::IpPattern),
|
||||
"repeat_offender" => Ok(Self::RepeatOffender),
|
||||
"frequency" => Ok(Self::Frequency),
|
||||
"multi_source_min" => Ok(Self::MultiSourceMin),
|
||||
"single_source_high" => Ok(Self::SingleSourceHigh),
|
||||
"fused_confidence_above" => Ok(Self::FusedConfidenceAbove),
|
||||
other => Err(SoarError::UnknownConditionType(other)),
|
||||
}
|
||||
}
|
||||
|
||||
@ -100,6 +100,11 @@ pub struct MLInferenceConfig {
|
||||
pub c2_threshold: f32,
|
||||
#[serde(default = "default_class_min_confidence")]
|
||||
pub class_min_confidence: f32,
|
||||
/// Multiplier applied to the confidence threshold before the aggregator
|
||||
/// fires an alert. The manifest can override this via
|
||||
/// `thresholds.alert_multiplier`.
|
||||
#[serde(default = "default_alert_threshold_multiplier")]
|
||||
pub alert_threshold_multiplier: f32,
|
||||
pub model_type: String,
|
||||
pub output_names: Vec<String>,
|
||||
pub ae_feature_weights: HashMap<String, f64>,
|
||||
@ -109,6 +114,10 @@ fn default_class_min_confidence() -> f32 {
|
||||
0.4
|
||||
}
|
||||
|
||||
fn default_alert_threshold_multiplier() -> f32 {
|
||||
1.2
|
||||
}
|
||||
|
||||
impl MLInferenceConfig {
|
||||
pub fn num_ae_features(&self) -> usize {
|
||||
self.ae_feature_names.len()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user