compliance-scanner-agent/compliance-agent/src/pipeline/cve.rs

use compliance_core::models::{CveAlert, CveSource, SbomEntry, VulnRef};
use compliance_core::CoreError;

pub struct CveScanner {
    http: reqwest::Client,
    #[allow(dead_code)]
    searxng_url: Option<String>,
    nvd_api_key: Option<String>,
}

impl CveScanner {
    pub fn new(
        http: reqwest::Client,
        searxng_url: Option<String>,
        nvd_api_key: Option<String>,
    ) -> Self {
        Self {
            http,
            searxng_url,
            nvd_api_key,
        }
    }

    #[tracing::instrument(skip_all)]
    pub async fn scan_dependencies(
        &self,
        repo_id: &str,
        entries: &mut [SbomEntry],
    ) -> Result<Vec<CveAlert>, CoreError> {
        tracing::info!("scanning {} SBOM entries for known CVEs", entries.len());
        let mut alerts = Vec::new();

        // Batch query OSV.dev
        let osv_results = self.query_osv_batch(entries).await?;
        for (idx, vulns) in osv_results.into_iter().enumerate() {
            if let Some(entry) = entries.get_mut(idx) {
                for vuln in &vulns {
                    entry.known_vulnerabilities.push(VulnRef {
                        id: vuln.id.clone(),
                        source: "osv".to_string(),
                        severity: vuln.severity.clone(),
                        url: Some(format!("https://osv.dev/vulnerability/{}", vuln.id)),
                    });

                    let mut alert = CveAlert::new(
                        vuln.id.clone(),
                        repo_id.to_string(),
                        entry.name.clone(),
                        entry.version.clone(),
                        CveSource::Osv,
                    );
                    alert.summary = vuln.summary.clone();
                    alerts.push(alert);
                }
            }
        }

        // Enrich with NVD CVSS scores
        for alert in &mut alerts {
            if let Ok(Some(cvss)) = self.query_nvd(&alert.cve_id).await {
                alert.cvss_score = Some(cvss);
            }
        }

        Ok(alerts)
    }

    async fn query_osv_batch(&self, entries: &[SbomEntry]) -> Result<Vec<Vec<OsvVuln>>, CoreError> {
        const OSV_BATCH_SIZE: usize = 500;

        let queries: Vec<_> = entries
            .iter()
            .filter_map(|e| {
                e.purl.as_ref().map(|purl| {
                    serde_json::json!({
                        "package": { "purl": purl }
                    })
                })
            })
            .collect();

        if queries.is_empty() {
            return Ok(Vec::new());
        }

        let mut all_vulns: Vec<Vec<OsvVuln>> = Vec::with_capacity(queries.len());

        for chunk in queries.chunks(OSV_BATCH_SIZE) {
            let body = serde_json::json!({ "queries": chunk });

            let resp = self
                .http
                .post("https://api.osv.dev/v1/querybatch")
                .json(&body)
                .send()
                .await
                .map_err(|e| {
                    tracing::warn!("OSV.dev API call failed: {e}");
                    CoreError::Http(format!("OSV.dev request failed: {e}"))
                })?;

            if !resp.status().is_success() {
                let status = resp.status();
                let body = resp.text().await.unwrap_or_default();
                tracing::warn!("OSV.dev returned {status}: {body}");
                // Push empty results for this chunk so indices stay aligned
                all_vulns.extend(std::iter::repeat_with(Vec::new).take(chunk.len()));
                continue;
            }

            let result: OsvBatchResponse = resp.json().await.map_err(|e| {
                tracing::warn!("failed to parse OSV.dev response: {e}");
                CoreError::Http(format!("Failed to parse OSV.dev response: {e}"))
            })?;

            let chunk_vulns = result.results.into_iter().map(|r| {
                r.vulns
                    .unwrap_or_default()
                    .into_iter()
                    .map(|v| OsvVuln {
                        id: v.id,
                        summary: v.summary,
                        severity: v.database_specific.and_then(|d| {
                            d.get("severity").and_then(|s| s.as_str()).map(String::from)
                        }),
                    })
                    .collect()
            });

            all_vulns.extend(chunk_vulns);
        }

        Ok(all_vulns)
    }

    async fn query_nvd(&self, cve_id: &str) -> Result<Option<f64>, CoreError> {
        if !cve_id.starts_with("CVE-") {
            return Ok(None);
        }

        let url = format!("https://services.nvd.nist.gov/rest/json/cves/2.0?cveId={cve_id}");
        let mut req = self.http.get(&url);

        if let Some(key) = &self.nvd_api_key {
            req = req.header("apiKey", key.as_str());
        }

        let resp = req
            .send()
            .await
            .map_err(|e| CoreError::Http(format!("NVD request failed: {e}")))?;

        if !resp.status().is_success() {
            return Ok(None);
        }

        let body: serde_json::Value = resp
            .json()
            .await
            .map_err(|e| CoreError::Http(format!("Failed to parse NVD response: {e}")))?;

        // Extract CVSS v3.1 base score
        let score = body["vulnerabilities"]
            .as_array()
            .and_then(|v| v.first())
            .and_then(|v| v["cve"]["metrics"]["cvssMetricV31"].as_array())
            .and_then(|m| m.first())
            .and_then(|m| m["cvssData"]["baseScore"].as_f64());

        Ok(score)
    }

    #[allow(dead_code)]
    pub async fn search_context(&self, cve_id: &str) -> Result<Vec<String>, CoreError> {
        let Some(searxng_url) = &self.searxng_url else {
            return Ok(Vec::new());
        };

        let url = format!(
            "{}/search?q={cve_id}&format=json&engines=duckduckgo",
            searxng_url.trim_end_matches('/')
        );
        let resp = self
            .http
            .get(&url)
            .send()
            .await
            .map_err(|e| CoreError::Http(format!("SearXNG request failed: {e}")))?;

        if !resp.status().is_success() {
            return Ok(Vec::new());
        }

        let body: serde_json::Value = resp.json().await.unwrap_or_default();
        let results = body["results"]
            .as_array()
            .map(|arr| {
                arr.iter()
                    .take(5)
                    .filter_map(|r| r["url"].as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_default();

        Ok(results)
    }
}

#[derive(serde::Deserialize)]
struct OsvBatchResponse {
    results: Vec<OsvBatchResult>,
}

#[derive(serde::Deserialize)]
struct OsvBatchResult {
    vulns: Option<Vec<OsvVulnEntry>>,
}

#[derive(serde::Deserialize)]
struct OsvVulnEntry {
    id: String,
    summary: Option<String>,
    database_specific: Option<serde_json::Value>,
}

struct OsvVuln {
    id: String,
    summary: Option<String>,
    severity: Option<String>,
}