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> { 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); } } }