Mantis/mantis/src/detection/ml/engine.rs

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use macros::log;
use tokio::sync::oneshot;
use tokio::time::interval;

use super::aggregator::AttackAggregator;
use super::config_loader::InferenceConfig;
use super::feature_extractor::FlowFeatures;
use super::flow_tracker::FlowTracker;
use super::inference::Inference;
use super::model_loader::MLModels;
use super::traffic_logger::TrafficLogger;
use crate::detection::fusion::FusionEngine;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::{EngineStats, InferenceStats};
use crate::utils::cpu_affinity::{num_cpus, set_cpu_affinity};
use crate::utils::packet_parser::parse_packet;

pub struct Engine {
    tracker: Arc<Mutex<FlowTracker>>,
    inference_pipeline: Arc<Inference>,
    aggregator: Arc<Mutex<AttackAggregator>>,
    fusion_engine: Arc<FusionEngine>,
    batch_size: usize,
    inference_interval_secs: u64,
    flow_timeout_us: u64,
    traffic_logger: Option<Arc<TrafficLogger>>,
    ml_cpu: Option<u32>,
}

impl Engine {
    pub fn new(
        models: Arc<MLModels>,
        config: Arc<InferenceConfig>,
        fusion_engine: Arc<FusionEngine>,
        max_flows: usize,
        batch_size: usize,
        interval_secs: u64,
        window_secs: u64,
        flow_timeout_us: u64,
        traffic_logger: Option<Arc<TrafficLogger>>,
        ml_cpu: Option<u32>,
    ) -> Self {
        let tracker = Arc::new(Mutex::new(FlowTracker::new(max_flows)));
        let inference_pipeline = Arc::new(Inference::new(models, config));

        let min_detections = ((window_secs / interval_secs) / 2).max(1) as usize;
        let aggregator = Arc::new(Mutex::new(AttackAggregator::new(window_secs, min_detections)));

        Self {
            tracker,
            inference_pipeline,
            aggregator,
            fusion_engine,
            batch_size,
            inference_interval_secs: interval_secs,
            flow_timeout_us,
            traffic_logger,
            ml_cpu,
        }
    }

    pub async fn run(self: Arc<Self>) -> oneshot::Sender<()> {
        let (shutdown_tx, shutdown_rx) = oneshot::channel();
        tokio::spawn(async move {
            self.run_inference_loop(shutdown_rx).await;
        });
        shutdown_tx
    }

    pub fn tracker(&self) -> &Arc<Mutex<FlowTracker>> {
        &self.tracker
    }

    async fn run_inference_loop(&self, mut shutdown_rx: oneshot::Receiver<()>) {
        let mut ticker = interval(Duration::from_secs(self.inference_interval_secs));

        loop {
            tokio::select! {
                _ = &mut shutdown_rx => break,
                _ = ticker.tick() => {}
            }

            let (total_flows, flows, active_ips) = {
                let Ok(mut t) = self.tracker.lock() else {
                    log!(MLError::TrackerLockPoisoned);
                    continue;
                };
                let total_flows = t.flow_count();
                let flows = t.drain_flows();
                t.cleanup_old_flows(self.flow_timeout_us);
                // Build active_ips from the drained flows, not the post-drain tracker
                // (which is always empty). This preserves per-src_ip inference buffers
                // across consecutive cycles so the LSTM window can fill up over time.
                // A src_ip absent from this cycle loses its buffer on the next cleanup.
                let active_ips: std::collections::HashSet<String> =
                    flows.iter().map(|f| f.flow_key.src_ip.clone()).collect();
                (total_flows, flows, active_ips)
            };

            log!(MLLog::FlowStats(total_flows, flows.len()));

            if flows.is_empty() {
                self.inference_pipeline.cleanup_buffers(&active_ips);
                continue;
            }

            if let Some(ref logger) = self.traffic_logger {
                let feature_names = FlowFeatures::all_feature_names_owned();
                for flow in &flows {
                    let features = FlowFeatures::extract(flow, &feature_names);
                    logger.log_row(features.to_csv_record());
                }
                continue;
            }

            let mut batch = flows[..flows.len().min(self.batch_size)].to_vec();
            batch.sort_by(|a, b| {
                a.flow_key
                    .src_ip
                    .cmp(&b.flow_key.src_ip)
                    .then_with(|| a.start_time_us.cmp(&b.start_time_us))
            });

            let start = Instant::now();
            let batch_len = batch.len();
            let pipeline = Arc::clone(&self.inference_pipeline);
            let ml_cpu = self.ml_cpu;
            let mut handle = tokio::task::spawn_blocking(move || {
                let cpu = ml_cpu
                    .map(|c| c as usize)
                    .unwrap_or_else(|| num_cpus().saturating_sub(1));
                set_cpu_affinity(cpu);
                pipeline.infer_batch(&batch)
            });

            let results = tokio::select! {
                _ = &mut shutdown_rx => {
                    handle.abort();
                    break;
                }
                res = &mut handle => res.unwrap_or_default(),
            };
            let elapsed_us = start.elapsed().as_micros() as u64;

            if results.is_empty() {
                // All flows are still in the warm-up window; no ONNX inference ran.
                self.inference_pipeline.cleanup_buffers(&active_ips);
                continue;
            }

            let stats = InferenceStats::from_results(&results, elapsed_us);
            log!(MLLog::InferenceCompleted(
                results.len(),
                batch_len,
                stats.malicious_flows,
                stats.benign_flows,
                (elapsed_us as f64 / 1000.0) as u32,
                stats.flows_per_second
            ));

            if let Ok(mut aggregator) = self.aggregator.lock() {
                let mut alerted_src_ips = std::collections::HashSet::new();
                for result in &results {
                    if result.is_attack {
                        // L1: full 5-tuple aggregation for persistent same-port attacks
                        let should_alert =
                            aggregator.should_alert(&result.flow_key_raw, result.ae_score, result.threshold);
                        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.fusion_engine.record_ml(result);
                            alerted_src_ips.insert(result.flow_key_raw.src_ip.clone());
                        }

                        // L2: src_ip-level scan/flood detection (skipped if L1 already fired)
                        if !alerted_src_ips.contains(&result.flow_key_raw.src_ip) {
                            if let Some(attack_type) = aggregator
                                .should_alert_src_ip(&result.flow_key_raw.src_ip, result.flow_key_raw.dst_port)
                            {
                                let mut l2_result = result.clone();
                                l2_result.attack_type = Some(attack_type.to_string());
                                log!(MLLog::ThreatDetected(
                                    format!("{:?}", l2_result.direction),
                                    l2_result.flow_key.clone(),
                                    attack_type.to_string(),
                                    l2_result.confidence,
                                    l2_result.ae_score,
                                ));
                                self.fusion_engine.record_ml(&l2_result);
                                alerted_src_ips.insert(result.flow_key_raw.src_ip.clone());
                            }
                        }
                    }
                }
                aggregator.cleanup();
            }

            self.inference_pipeline.cleanup_buffers(&active_ips);
        }
    }

    pub fn process_packet(&self, packet_data: &[u8], is_ingress: bool) {
        match parse_packet(packet_data) {
            Some((packet_info, payload_start)) => {
                let payload = packet_data.get(payload_start..).unwrap_or(&[]);
                if let Ok(mut t) = self.tracker.lock() {
                    t.process_packet(packet_info, is_ingress, payload);
                } else {
                    log!(MLError::TrackerLockPoisoned);
                }
            }
            None => log!(MLLog::ParsePacketFailed(packet_data.len())),
        }
    }

    pub fn get_stats(&self) -> EngineStats {
        let active_flows = match self.tracker.lock() {
            Ok(t) => t.flow_count(),
            Err(_) => {
                log!(MLError::TrackerLockPoisoned);
                0
            }
        };
        EngineStats { active_flows }
    }
}