WTFnet/crates/wtfnet-dnsleak/src/report.rs

use crate::policy::PolicySummary;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::net::IpAddr;
use wtfnet_platform::{FlowOwner, FlowOwnerConfidence, FlowProtocol};

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LeakTransport {
    Udp53,
    Tcp53,
    Dot,
    Doh,
    Unknown,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[serde(rename_all = "lowercase")]
pub enum LeakType {
    A,
    B,
    C,
    D,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum RouteClass {
    Loopback,
    Tunnel,
    Physical,
    Unknown,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Severity {
    P0,
    P1,
    P2,
    P3,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnrichedEvent {
    pub timestamp_ms: u128,
    pub proto: FlowProtocol,
    pub src_ip: IpAddr,
    pub src_port: u16,
    pub dst_ip: IpAddr,
    pub dst_port: u16,
    pub iface_name: Option<String>,
    pub transport: LeakTransport,
    pub qname: Option<String>,
    pub qtype: Option<String>,
    pub rcode: Option<String>,
    pub is_response: bool,
    pub answer_ips: Vec<IpAddr>,
    pub route_class: RouteClass,
    pub owner: Option<FlowOwner>,
    pub owner_confidence: FlowOwnerConfidence,
    pub owner_failure: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeakEvent {
    pub timestamp_ms: u128,
    pub transport: LeakTransport,
    pub qname: Option<String>,
    pub qtype: Option<String>,
    pub rcode: Option<String>,
    pub iface_name: Option<String>,
    pub route_class: RouteClass,
    pub dst_ip: String,
    pub dst_port: u16,
    pub pid: Option<u32>,
    pub ppid: Option<u32>,
    pub process_name: Option<String>,
    pub process_path: Option<String>,
    pub attribution_confidence: FlowOwnerConfidence,
    pub attribution_failure: Option<String>,
    pub leak_type: LeakType,
    pub severity: Severity,
    pub policy_rule_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeakTypeCount {
    pub leak_type: LeakType,
    pub count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SummaryItem {
    pub key: String,
    pub count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeakSummary {
    pub total: usize,
    pub by_type: Vec<LeakTypeCount>,
    pub top_processes: Vec<SummaryItem>,
    pub top_destinations: Vec<SummaryItem>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeakReport {
    pub duration_ms: u64,
    pub policy: PolicySummary,
    pub summary: LeakSummary,
    pub events: Vec<LeakEvent>,
}

impl LeakEvent {
    pub fn from_decision(event: EnrichedEvent, decision: crate::rules::LeakDecision) -> Self {
        let (pid, ppid, process_name, process_path) = event
            .owner
            .as_ref()
            .map(|owner| {
                (
                    owner.pid,
                    owner.ppid,
                    owner.process_name.clone(),
                    owner.process_path.clone(),
                )
            })
            .unwrap_or((None, None, None, None));

        LeakEvent {
            timestamp_ms: event.timestamp_ms,
            transport: event.transport,
            qname: event.qname,
            qtype: event.qtype,
            rcode: event.rcode,
            iface_name: event.iface_name,
            route_class: event.route_class,
            dst_ip: event.dst_ip.to_string(),
            dst_port: event.dst_port,
            pid,
            ppid,
            process_name,
            process_path,
            attribution_confidence: event.owner_confidence,
            attribution_failure: event.owner_failure,
            leak_type: decision.leak_type,
            severity: decision.severity,
            policy_rule_id: decision.policy_rule_id,
        }
    }
}

impl LeakSummary {
    pub fn from_events(events: &[LeakEvent]) -> Self {
        let total = events.len();
        let mut by_type_map: HashMap<LeakType, usize> = HashMap::new();
        let mut process_map: BTreeMap<String, usize> = BTreeMap::new();
        let mut dest_map: BTreeMap<String, usize> = BTreeMap::new();

        for event in events {
            *by_type_map.entry(event.leak_type).or_insert(0) += 1;
            if let Some(name) = event.process_name.as_ref() {
                *process_map.entry(name.clone()).or_insert(0) += 1;
            }
            let dst_key = format!("{}:{}", event.dst_ip, event.dst_port);
            *dest_map.entry(dst_key).or_insert(0) += 1;
        }

        let mut by_type = by_type_map
            .into_iter()
            .map(|(leak_type, count)| LeakTypeCount { leak_type, count })
            .collect::<Vec<_>>();
        by_type.sort_by(|a, b| a.leak_type.cmp(&b.leak_type));
        let top_processes = top_items(process_map, 5);
        let top_destinations = top_items(dest_map, 5);

        LeakSummary {
            total,
            by_type,
            top_processes,
            top_destinations,
        }
    }
}

fn top_items(map: BTreeMap<String, usize>, limit: usize) -> Vec<SummaryItem> {
    let mut items = map
        .into_iter()
        .map(|(key, count)| SummaryItem { key, count })
        .collect::<Vec<_>>();
    items.sort_by(|a, b| b.count.cmp(&a.count).then_with(|| a.key.cmp(&b.key)));
    items.truncate(limit);
    items
}