compliance-scanner-agent/compliance-agent/src/controls/scan_triage.rs

//! Scan-pipeline integration for control triage.
//!
//! After the deterministic tools have produced findings, this stamps each finding
//! with the compliance control(s) it's evidence for and marks control-level false
//! positives — using the ingested OSCAL catalog for control text, the
//! `control-map` LUT for the finding→control link, and the grounded LLM judge to
//! confirm. Skipped entirely unless breakpilot is configured.

use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

use compliance_core::control_check::{CandidateRegion, ControlCheckSpec};
use compliance_core::models::finding::{Finding, FindingStatus, Severity};
use compliance_core::models::onboarding::ComplianceFramework;
use compliance_core::AgentConfig;
use control_map::ControlMap;

use super::surface;
use super::{
    ControlIndex, ControlTriage, GroundedControlChecker, LlmControlJudge, OscalControlsProvider,
    SemanticControlChecker, TriageOutcome,
};
use crate::llm::LlmClient;

/// Nearest master controls judged per code region in the semantic pass.
const SEMANTIC_TOP_K: usize = 5;

/// Lines of context to read on each side of a finding's line.
const REGION_WINDOW: usize = 6;

/// Triage every finding in `findings` against the CRA control map: stamp
/// `control_refs` on confirmed findings and flag control false positives. Returns
/// the number of findings tagged with at least one control.
pub async fn triage_repo_findings(
    config: &AgentConfig,
    llm: Arc<LlmClient>,
    repo_path: &Path,
    findings: &mut [Finding],
) -> usize {
    let Some(base_url) = config.breakpilot.base_url.clone() else {
        return 0; // control triage is opt-in via BREAKPILOT_BASE_URL
    };
    let provider = OscalControlsProvider::new(
        reqwest::Client::new(),
        base_url,
        config.breakpilot.token.clone(),
        &config.breakpilot.snapshot_dir,
    );
    let specs = build_specs(&provider).await;
    if specs.is_empty() {
        return 0;
    }
    let map = match ControlMap::cra() {
        Ok(m) => m,
        Err(e) => {
            tracing::warn!(error = %e, "control map failed to load; skipping control triage");
            return 0;
        }
    };
    let triage = ControlTriage::new(LlmControlJudge::new(llm), map, specs);

    let mut tagged = 0;
    for finding in findings.iter_mut() {
        let (Some(file), Some(line)) = (finding.file_path.clone(), finding.line_number) else {
            continue;
        };
        let Some(region) = fetch_region(repo_path, &file, line) else {
            continue;
        };
        match triage.triage(finding, &region).await {
            TriageOutcome::Confirmed(controls) => {
                finding.control_refs = controls;
                tagged += 1;
            }
            TriageOutcome::FalsePositive => {
                finding.status = FindingStatus::FalsePositive;
                finding.triage_action = Some("control_false_positive".to_string());
            }
            TriageOutcome::Unmapped => {}
        }
    }
    tagged
}

/// Build the control requirement specs (by id) from the ingested OSCAL catalog.
async fn build_specs(provider: &OscalControlsProvider) -> HashMap<String, ControlCheckSpec> {
    let mut specs = HashMap::new();
    match provider.load(ComplianceFramework::Cra).await {
        Ok(doc) => {
            for control in doc.to_controls() {
                specs.insert(
                    control.id.clone(),
                    ControlCheckSpec {
                        control_id: control.id,
                        title: control.title,
                        requirement: control.text,
                        default_cwe: None,
                        severity: Severity::Medium,
                    },
                );
            }
        }
        Err(e) => tracing::warn!(error = %e, "could not load control catalog for triage"),
    }
    specs
}

/// Absence-based control pass (the grounded half of the hybrid coverage): for each
/// control with a [`surface`] definition, deterministically retrieve the code
/// surfaces it governs (login routes, logging setup, update/download code) and have
/// the grounded judge decide whether the control holds there. Returns net-new
/// findings, each already tagged with its control and grounded to a real snippet.
///
/// The orchestrator runs this when `breakpilot.grounded_control_checks` is set
/// (on by default). Validated live; it covers the 8 absence-based CRA controls
/// (the judge decides presence/absence, grounded to a real snippet).
pub async fn grounded_surface_findings(
    config: &AgentConfig,
    llm: Arc<LlmClient>,
    repo_path: &Path,
    repo_id: &str,
) -> Vec<Finding> {
    let Some(base_url) = config.breakpilot.base_url.clone() else {
        return Vec::new();
    };
    let provider = OscalControlsProvider::new(
        reqwest::Client::new(),
        base_url,
        config.breakpilot.token.clone(),
        &config.breakpilot.snapshot_dir,
    );
    let specs = build_specs(&provider).await;
    if specs.is_empty() {
        return Vec::new();
    }
    let checker = GroundedControlChecker::new(LlmControlJudge::new(llm));

    let mut out = Vec::new();
    for surf in surface::SURFACES {
        let Some(spec) = specs.get(surf.control_id) else {
            continue; // catalog doesn't carry this control
        };
        let regions = surface::retrieve(repo_path, surf.terms);
        if regions.is_empty() {
            continue;
        }
        out.extend(checker.check(spec, &regions, repo_id).await);
    }
    out
}

/// Read a window of lines around `line` (1-based) from `repo_path/file`.
fn fetch_region(repo_path: &Path, file: &str, line: u32) -> Option<CandidateRegion> {
    let content = std::fs::read_to_string(repo_path.join(file)).ok()?;
    let lines: Vec<&str> = content.lines().collect();
    if lines.is_empty() {
        return None;
    }
    let center = (line.saturating_sub(1) as usize).min(lines.len() - 1);
    let start = center.saturating_sub(REGION_WINDOW);
    let end = (center + REGION_WINDOW + 1).min(lines.len());
    Some(CandidateRegion {
        file: file.to_string(),
        start_line: (start as u32) + 1,
        content: lines[start..end].join("\n"),
    })
}

/// Master-controls **semantic** pass: for each finding's code region, retrieve the
/// top-K nearest master controls by embedding, have the grounded judge confirm,
/// and stamp the confirmed control ids onto the finding — the scale path for the
/// ~13.6k master-control corpus (which has no CWE to LUT on). Returns the number
/// of findings that gained a master-control ref.
///
/// The orchestrator runs this when `breakpilot.semantic_mapping` is set (on by
/// default). The control embedding index is built once and cached to
/// `snapshot_dir` keyed by corpus hash ([`ControlIndex::load_or_build`]), so only
/// the first scan after a catalog change pays the embedding cost.
pub async fn semantic_stamp_findings(
    config: &AgentConfig,
    llm: Arc<LlmClient>,
    repo_path: &Path,
    findings: &mut [Finding],
) -> usize {
    let Some(base_url) = config.breakpilot.base_url.clone() else {
        return 0;
    };
    let provider = OscalControlsProvider::new(
        reqwest::Client::new(),
        base_url,
        config.breakpilot.token.clone(),
        &config.breakpilot.snapshot_dir,
    );
    let doc = match provider.load_master_controls().await {
        Ok(d) => d,
        Err(e) => {
            tracing::warn!(error = %e, "master-controls catalog unavailable; skipping semantic pass");
            return 0;
        }
    };
    let specs: Vec<ControlCheckSpec> = doc
        .to_controls()
        .into_iter()
        .map(|c| ControlCheckSpec {
            control_id: c.id,
            title: c.title,
            requirement: c.text,
            default_cwe: None,
            severity: Severity::Medium,
        })
        .collect();
    let cache_path =
        Path::new(&config.breakpilot.snapshot_dir).join("control-index-master-controls.json");
    let index = match ControlIndex::load_or_build(&llm, specs, &cache_path).await {
        Ok(i) if !i.is_empty() => i,
        Ok(_) => return 0,
        Err(e) => {
            tracing::warn!(error = %e, "failed to embed master-controls corpus");
            return 0;
        }
    };
    let checker = SemanticControlChecker::new(LlmControlJudge::new(llm.clone()));

    let mut tagged = 0;
    for finding in findings.iter_mut() {
        if finding.status == FindingStatus::FalsePositive {
            continue;
        }
        let (Some(file), Some(line)) = (finding.file_path.clone(), finding.line_number) else {
            continue;
        };
        let Some(region) = fetch_region(repo_path, &file, line) else {
            continue;
        };
        // Retrieve on the finding's intent + the code, not the region alone: two
        // findings in one file share overlapping windows and otherwise embed alike,
        // collapsing onto the same controls. The finding's title/description carry
        // the discriminating signal (e.g. "brute-force protection" vs "weak hash").
        // The raw `region` still goes to the judge for snippet grounding.
        let query = format!(
            "{}\n{}\n\n{}",
            finding.title, finding.description, region.content
        );
        let query_emb = match llm.embed(vec![query]).await {
            Ok(mut embs) => match embs.pop() {
                Some(v) => v,
                None => continue,
            },
            Err(e) => {
                tracing::warn!(error = %e, "query embed failed; skipping finding");
                continue;
            }
        };
        let confirmed = checker
            .check(
                &index,
                &region,
                &query_emb,
                SEMANTIC_TOP_K,
                &finding.repo_id,
            )
            .await;
        let before = finding.control_refs.len();
        for f in confirmed {
            for cref in f.control_refs {
                if !finding.control_refs.contains(&cref) {
                    finding.control_refs.push(cref);
                }
            }
        }
        if finding.control_refs.len() > before {
            tagged += 1;
        }
    }
    tagged
}

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

    #[test]
    fn fetch_region_windows_around_the_line() {
        let dir = std::env::temp_dir().join(format!("triage-region-{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&dir).unwrap();
        let file = "a.py";
        std::fs::write(dir.join(file), "l1\nl2\nl3\nSECRET=1\nl5\nl6\n").unwrap();
        let r = fetch_region(&dir, file, 4).unwrap();
        assert!(r.content.contains("SECRET=1"));
        assert_eq!(r.start_line, 1); // window clamps to file start
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn fetch_region_missing_file_is_none() {
        assert!(fetch_region(Path::new("/nonexistent"), "nope.py", 1).is_none());
    }
}