acaua/crates/acaua-core/src/fusion.rs

use crate::error::Result;
use crate::pillar4::{self, FusionInput, HostScore};
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Alert {
    pub id: String,
    pub pillar: u8,
    pub score: f64,
    pub summary: String,
    pub evidence: serde_json::Value,
    pub created_at: String,
}

pub fn alert_orig_h(alert: &Alert) -> String {
    if let Some(h) = alert.evidence.get("orig_h").and_then(|v| v.as_str()) {
        if !h.is_empty() {
            return h.to_string();
        }
    }
    if let Some(rest) = alert.summary.strip_prefix("Beaconing detected: ") {
        if let Some(h) = rest.split(':').next() {
            return h.to_string();
        }
    }
    if let Some(rest) = alert.summary.strip_prefix("Persistent session: ") {
        if let Some(h) = rest.split(" -> ").next() {
            return h.to_string();
        }
    }
    if let Some(rest) = alert
        .summary
        .strip_prefix("High-information DNS subdomains: ")
    {
        if let Some(h) = rest.split(" -> ").next() {
            return h.to_string();
        }
    }
    String::new()
}

pub fn fuse(alerts: &[Alert]) -> Result<Vec<HostScore>> {
    let mut hosts_set = BTreeSet::new();
    for alert in alerts {
        let h = alert_orig_h(alert);
        if !h.is_empty() {
            hosts_set.insert(h);
        }
    }

    let mut host_scores = Vec::new();

    for orig_h in hosts_set {
        let mut m1 = 0.0_f64;
        let mut alert_id_p1 = None;
        let mut m2 = 0.0_f64;
        let mut alert_id_p2 = None;
        let mut m3 = 0.0_f64;
        let mut alert_id_p3 = None;

        for alert in alerts {
            if alert_orig_h(alert) == orig_h {
                match alert.pillar {
                    1 => {
                        if alert.score > m1 || (alert.score == m1 && alert_id_p1.is_none()) {
                            m1 = alert.score;
                            alert_id_p1 = Some(alert.id.clone());
                        } else if alert.score == m1 {
                            if let Some(prev) = &alert_id_p1 {
                                if alert.id < *prev {
                                    alert_id_p1 = Some(alert.id.clone());
                                }
                            }
                        }
                    }
                    2 => {
                        if alert.score > m2 || (alert.score == m2 && alert_id_p2.is_none()) {
                            m2 = alert.score;
                            alert_id_p2 = Some(alert.id.clone());
                        } else if alert.score == m2 {
                            if let Some(prev) = &alert_id_p2 {
                                if alert.id < *prev {
                                    alert_id_p2 = Some(alert.id.clone());
                                }
                            }
                        }
                    }
                    3 => {
                        if alert.score > m3 || (alert.score == m3 && alert_id_p3.is_none()) {
                            m3 = alert.score;
                            alert_id_p3 = Some(alert.id.clone());
                        } else if alert.score == m3 {
                            if let Some(prev) = &alert_id_p3 {
                                if alert.id < *prev {
                                    alert_id_p3 = Some(alert.id.clone());
                                }
                            }
                        }
                    }
                    _ => {}
                }
            }
        }

        let input = FusionInput {
            m1,
            m2,
            m3,
            alert_id_p1,
            alert_id_p2,
            alert_id_p3,
        };

        let hs = pillar4::score(&orig_h, &input)?;
        host_scores.push(hs);
    }

    host_scores.sort_by(|a, b| {
        b.score
            .total_cmp(&a.score)
            .then_with(|| a.orig_h.cmp(&b.orig_h))
    });

    Ok(host_scores)
}

#[cfg(test)]
mod tests {
    use super::*;

    struct SplitMix64 {
        state: u64,
    }

    impl SplitMix64 {
        fn new(seed: u64) -> Self {
            Self { state: seed }
        }
        fn next_u64(&mut self) -> u64 {
            self.state = self.state.wrapping_add(0x9e3779b97f4a7c15);
            let mut z = self.state;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
            z ^ (z >> 31)
        }
    }

    #[test]
    fn test_permutation_invariance_and_evidence() {
        let alerts = vec![
            Alert {
                id: "alert_p1_a".to_string(),
                pillar: 1,
                score: 0.85,
                summary: "Beaconing detected: 10.0.0.1:any -> 1.1.1.1:443 (period ~14400s, score 0.850)".to_string(),
                evidence: serde_json::json!({ "orig_h": "10.0.0.1" }),
                created_at: "2026-10-31T12:00:00Z".to_string(),
            },
            Alert {
                id: "alert_p2_a".to_string(),
                pillar: 2,
                score: 0.60,
                summary: "Persistent session: 10.0.0.1 -> 1.1.1.1:443 (6.0h, score 0.600)".to_string(),
                evidence: serde_json::json!({ "orig_h": "10.0.0.1" }),
                created_at: "2026-10-31T12:00:00Z".to_string(),
            },
            Alert {
                id: "alert_p3_a".to_string(),
                pillar: 3,
                score: 0.50,
                summary: "High-information DNS subdomains: 10.0.0.1 -> *.example.com (50000 bits/h, 100 unique, score 0.500)".to_string(),
                evidence: serde_json::json!({ "orig_h": "10.0.0.1" }),
                created_at: "2026-10-31T12:00:00Z".to_string(),
            },
            Alert {
                id: "alert_p1_b".to_string(),
                pillar: 1,
                score: 0.45,
                summary: "Beaconing detected: 10.0.0.2:any -> 2.2.2.2:80 (period ~900s, score 0.450)".to_string(),
                evidence: serde_json::json!({ "orig_h": "10.0.0.2" }),
                created_at: "2026-10-31T12:00:00Z".to_string(),
            },
        ];

        let reference = fuse(&alerts).unwrap();
        let ref_json = serde_json::to_vec(&reference).unwrap();

        // Check that each term.alert_id exists in original alerts
        for hs in &reference {
            for term in &hs.terms {
                if let Some(id) = &term.alert_id {
                    assert!(alerts.iter().any(|a| &a.id == id));
                }
            }
        }

        // Test 10 random permutations
        let mut rng = SplitMix64::new(42);
        for _ in 0..10 {
            let mut permuted = alerts.clone();
            for i in (1..permuted.len()).rev() {
                let j = (rng.next_u64() as usize) % (i + 1);
                permuted.swap(i, j);
            }
            let fused = fuse(&permuted).unwrap();
            let fused_json = serde_json::to_vec(&fused).unwrap();
            assert_eq!(ref_json, fused_json);
        }
    }
}