compliance-scanner-agent/compliance-dashboard/src/components/attack_chain/helpers.rs

use std::collections::{HashMap, VecDeque};

/// Get category CSS class from tool name
pub(crate) fn tool_category(name: &str) -> &'static str {
    let lower = name.to_lowercase();
    if lower.contains("recon") {
        return "recon";
    }
    if lower.contains("openapi") || lower.contains("api") || lower.contains("swagger") {
        return "api";
    }
    if lower.contains("header") {
        return "headers";
    }
    if lower.contains("csp") {
        return "csp";
    }
    if lower.contains("cookie") {
        return "cookies";
    }
    if lower.contains("log") || lower.contains("console") {
        return "logs";
    }
    if lower.contains("rate") || lower.contains("limit") {
        return "ratelimit";
    }
    if lower.contains("cors") {
        return "cors";
    }
    if lower.contains("tls") || lower.contains("ssl") {
        return "tls";
    }
    if lower.contains("redirect") {
        return "redirect";
    }
    if lower.contains("dns")
        || lower.contains("dmarc")
        || lower.contains("email")
        || lower.contains("spf")
    {
        return "email";
    }
    if lower.contains("auth")
        || lower.contains("jwt")
        || lower.contains("token")
        || lower.contains("session")
    {
        return "auth";
    }
    if lower.contains("xss") {
        return "xss";
    }
    if lower.contains("sql") || lower.contains("sqli") {
        return "sqli";
    }
    if lower.contains("ssrf") {
        return "ssrf";
    }
    if lower.contains("idor") {
        return "idor";
    }
    if lower.contains("fuzz") {
        return "fuzzer";
    }
    if lower.contains("cve") || lower.contains("exploit") {
        return "cve";
    }
    "default"
}

/// Get emoji icon from tool category
pub(crate) fn tool_emoji(cat: &str) -> &'static str {
    match cat {
        "recon" => "\u{1F50D}",
        "api" => "\u{1F517}",
        "headers" => "\u{1F6E1}",
        "csp" => "\u{1F6A7}",
        "cookies" => "\u{1F36A}",
        "logs" => "\u{1F4DD}",
        "ratelimit" => "\u{23F1}",
        "cors" => "\u{1F30D}",
        "tls" => "\u{1F510}",
        "redirect" => "\u{21AA}",
        "email" => "\u{1F4E7}",
        "auth" => "\u{1F512}",
        "xss" => "\u{26A1}",
        "sqli" => "\u{1F489}",
        "ssrf" => "\u{1F310}",
        "idor" => "\u{1F511}",
        "fuzzer" => "\u{1F9EA}",
        "cve" => "\u{1F4A3}",
        _ => "\u{1F527}",
    }
}

/// Compute display label for category
pub(crate) fn cat_label(cat: &str) -> &'static str {
    match cat {
        "recon" => "Recon",
        "api" => "API",
        "headers" => "Headers",
        "csp" => "CSP",
        "cookies" => "Cookies",
        "logs" => "Logs",
        "ratelimit" => "Rate Limit",
        "cors" => "CORS",
        "tls" => "TLS",
        "redirect" => "Redirect",
        "email" => "Email/DNS",
        "auth" => "Auth",
        "xss" => "XSS",
        "sqli" => "SQLi",
        "ssrf" => "SSRF",
        "idor" => "IDOR",
        "fuzzer" => "Fuzzer",
        "cve" => "CVE",
        _ => "Other",
    }
}

/// Maximum number of display phases — deeper iterations are merged into the last.
const MAX_PHASES: usize = 8;

/// Phase name heuristic based on phase index (not raw BFS depth)
pub(crate) fn phase_name(phase_idx: usize) -> &'static str {
    match phase_idx {
        0 => "Reconnaissance",
        1 => "Analysis",
        2 => "Boundary Testing",
        3 => "Injection & Exploitation",
        4 => "Authentication Testing",
        5 => "Validation",
        6 => "Deep Scan",
        _ => "Final",
    }
}

/// Short label for phase rail
pub(crate) fn phase_short_name(phase_idx: usize) -> &'static str {
    match phase_idx {
        0 => "Recon",
        1 => "Analysis",
        2 => "Boundary",
        3 => "Exploit",
        4 => "Auth",
        5 => "Validate",
        6 => "Deep",
        _ => "Final",
    }
}

/// Compute BFS phases from attack chain nodes
pub(crate) fn compute_phases(steps: &[serde_json::Value]) -> Vec<Vec<usize>> {
    let node_ids: Vec<String> = steps
        .iter()
        .map(|s| {
            s.get("node_id")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string()
        })
        .collect();

    let id_to_idx: HashMap<String, usize> = node_ids
        .iter()
        .enumerate()
        .map(|(i, id)| (id.clone(), i))
        .collect();

    // Compute depth via BFS
    let mut depths = vec![usize::MAX; steps.len()];
    let mut queue = VecDeque::new();

    // Root nodes: those with no parents or parents not in the set
    for (i, step) in steps.iter().enumerate() {
        let parents = step
            .get("parent_node_ids")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|p| p.as_str())
                    .filter(|p| id_to_idx.contains_key(*p))
                    .count()
            })
            .unwrap_or(0);
        if parents == 0 {
            depths[i] = 0;
            queue.push_back(i);
        }
    }

    // BFS to compute min depth
    while let Some(idx) = queue.pop_front() {
        let current_depth = depths[idx];
        let node_id = &node_ids[idx];
        // Find children: nodes that list this node as a parent
        for (j, step) in steps.iter().enumerate() {
            if depths[j] <= current_depth + 1 {
                continue;
            }
            let is_child = step
                .get("parent_node_ids")
                .and_then(|v| v.as_array())
                .map(|arr| arr.iter().any(|p| p.as_str() == Some(node_id.as_str())))
                .unwrap_or(false);
            if is_child {
                depths[j] = current_depth + 1;
                queue.push_back(j);
            }
        }
    }

    // Handle unreachable nodes
    for d in depths.iter_mut() {
        if *d == usize::MAX {
            *d = 0;
        }
    }

    // Cap depths at MAX_PHASES - 1 so deeper iterations merge into the last phase
    for d in depths.iter_mut() {
        if *d >= MAX_PHASES {
            *d = MAX_PHASES - 1;
        }
    }

    // Group by (capped) depth
    let max_depth = depths.iter().copied().max().unwrap_or(0);
    let mut phases: Vec<Vec<usize>> = Vec::new();
    for d in 0..=max_depth {
        let indices: Vec<usize> = depths
            .iter()
            .enumerate()
            .filter(|(_, &dep)| dep == d)
            .map(|(i, _)| i)
            .collect();
        if !indices.is_empty() {
            phases.push(indices);
        }
    }
    phases
}

/// Format BSON datetime to readable string
pub(crate) fn format_bson_time(val: &serde_json::Value) -> String {
    // Handle BSON {"$date":{"$numberLong":"..."}}
    if let Some(date_obj) = val.get("$date") {
        if let Some(ms_str) = date_obj.get("$numberLong").and_then(|v| v.as_str()) {
            if let Ok(ms) = ms_str.parse::<i64>() {
                let secs = ms / 1000;
                let h = (secs / 3600) % 24;
                let m = (secs / 60) % 60;
                let s = secs % 60;
                return format!("{h:02}:{m:02}:{s:02}");
            }
        }
        // Handle {"$date": "2025-..."}
        if let Some(s) = date_obj.as_str() {
            return s.to_string();
        }
    }
    // Handle plain string
    if let Some(s) = val.as_str() {
        return s.to_string();
    }
    String::new()
}

/// Compute duration string from started_at and completed_at
pub(crate) fn compute_duration(step: &serde_json::Value) -> String {
    let extract_ms = |val: &serde_json::Value| -> Option<i64> {
        val.get("$date")?
            .get("$numberLong")?
            .as_str()?
            .parse::<i64>()
            .ok()
    };

    let started = step.get("started_at").and_then(extract_ms);
    let completed = step.get("completed_at").and_then(extract_ms);

    match (started, completed) {
        (Some(s), Some(c)) => {
            let diff_ms = c - s;
            if diff_ms < 1000 {
                format!("{}ms", diff_ms)
            } else {
                format!("{:.1}s", diff_ms as f64 / 1000.0)
            }
        }
        _ => String::new(),
    }
}