compliance-scanner-agent/compliance-agent/src/pentest/report/html/attack_chain.rs

use super::html_escape;
use compliance_core::models::pentest::AttackChainNode;

pub(super) fn attack_chain(chain: &[AttackChainNode]) -> String {
    let chain_section = if chain.is_empty() {
        r#"<p style="color: var(--text-muted);">No attack chain steps recorded.</p>"#.to_string()
    } else {
        build_chain_html(chain)
    };

    format!(
        r##"<!-- ═══════════════ 4. ATTACK CHAIN ═══════════════ -->
<div class="page-break"></div>
<h2><span class="section-num">4.</span> Attack Chain Timeline</h2>

<p>
    The following sequence shows each tool invocation made by the AI orchestrator during the assessment,
    grouped by phase. Each step includes the tool's name, execution status, and the AI's reasoning
    for choosing that action.
</p>

<div style="margin-top: 16px;">
    {chain_section}
</div>"##
    )
}

fn build_chain_html(chain: &[AttackChainNode]) -> String {
    let mut chain_html = String::new();

    // Compute phases via BFS from root nodes
    let mut phase_map: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
    let mut queue: std::collections::VecDeque<String> = std::collections::VecDeque::new();

    for node in chain {
        if node.parent_node_ids.is_empty() {
            let nid = node.node_id.clone();
            if !nid.is_empty() {
                phase_map.insert(nid.clone(), 0);
                queue.push_back(nid);
            }
        }
    }

    while let Some(nid) = queue.pop_front() {
        let parent_phase = phase_map.get(&nid).copied().unwrap_or(0);
        for node in chain {
            if node.parent_node_ids.contains(&nid) {
                let child_id = node.node_id.clone();
                if !child_id.is_empty() && !phase_map.contains_key(&child_id) {
                    phase_map.insert(child_id.clone(), parent_phase + 1);
                    queue.push_back(child_id);
                }
            }
        }
    }

    // Assign phase 0 to any unassigned nodes
    for node in chain {
        let nid = node.node_id.clone();
        if !nid.is_empty() && !phase_map.contains_key(&nid) {
            phase_map.insert(nid, 0);
        }
    }

    // Group nodes by phase
    let max_phase = phase_map.values().copied().max().unwrap_or(0);
    let phase_labels = [
        "Reconnaissance",
        "Enumeration",
        "Exploitation",
        "Validation",
        "Post-Exploitation",
    ];

    for phase_idx in 0..=max_phase {
        let phase_nodes: Vec<&AttackChainNode> = chain
            .iter()
            .filter(|n| {
                let nid = n.node_id.clone();
                phase_map.get(&nid).copied().unwrap_or(0) == phase_idx
            })
            .collect();

        if phase_nodes.is_empty() {
            continue;
        }

        let label = if phase_idx < phase_labels.len() {
            phase_labels[phase_idx]
        } else {
            "Additional Testing"
        };

        chain_html.push_str(&format!(
            r#"<div class="phase-block">
                    <div class="phase-header">
                        <span class="phase-num">Phase {}</span>
                        <span class="phase-label">{}</span>
                        <span class="phase-count">{} step{}</span>
                    </div>
                    <div class="phase-steps">"#,
            phase_idx + 1,
            label,
            phase_nodes.len(),
            if phase_nodes.len() == 1 { "" } else { "s" },
        ));

        for (i, node) in phase_nodes.iter().enumerate() {
            let status_label = format!("{:?}", node.status);
            let status_class = match status_label.to_lowercase().as_str() {
                "completed" => "step-completed",
                "failed" => "step-failed",
                _ => "step-running",
            };
            let findings_badge = if !node.findings_produced.is_empty() {
                format!(
                    r#"<span class="step-findings">{} finding{}</span>"#,
                    node.findings_produced.len(),
                    if node.findings_produced.len() == 1 {
                        ""
                    } else {
                        "s"
                    },
                )
            } else {
                String::new()
            };
            let risk_badge = node
                .risk_score
                .map(|r| {
                    let risk_class = if r >= 70 {
                        "risk-high"
                    } else if r >= 40 {
                        "risk-med"
                    } else {
                        "risk-low"
                    };
                    format!(r#"<span class="step-risk {risk_class}">Risk: {r}</span>"#)
                })
                .unwrap_or_default();

            let reasoning_html = if node.llm_reasoning.is_empty() {
                String::new()
            } else {
                format!(
                    r#"<div class="step-reasoning">{}</div>"#,
                    html_escape(&node.llm_reasoning)
                )
            };

            // Render inline screenshot if this is a browser screenshot action
            let screenshot_html = if node.tool_name == "browser" {
                node.tool_output
                    .as_ref()
                    .and_then(|out| out.get("screenshot_base64"))
                    .and_then(|v| v.as_str())
                    .filter(|s| !s.is_empty())
                    .map(|b64| {
                        format!(
                            r#"<div class="step-screenshot"><img src="data:image/png;base64,{b64}" alt="Browser screenshot" style="max-width:100%;border:1px solid #e2e8f0;border-radius:6px;margin-top:8px;"/></div>"#
                        )
                    })
                    .unwrap_or_default()
            } else {
                String::new()
            };

            chain_html.push_str(&format!(
                r#"<div class="step-row">
                        <div class="step-num">{num}</div>
                        <div class="step-connector"></div>
                        <div class="step-content">
                            <div class="step-header">
                                <span class="step-tool">{tool_name}</span>
                                <span class="step-status {status_class}">{status_label}</span>
                                {findings_badge}
                                {risk_badge}
                            </div>
                            {reasoning_html}
                            {screenshot_html}
                        </div>
                    </div>"#,
                num = i + 1,
                tool_name = html_escape(&node.tool_name),
            ));
        }

        chain_html.push_str("</div></div>");
    }

    chain_html
}