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

//! Pull + snapshot [`ControlsProvider`] backed by breakpilot-compliance's OSCAL
//! catalog export.
//!
//! Fetches `GET {base}/api/compliance/v1/oscal/catalog?framework=<fw>`, snapshots
//! the exact bytes to disk (so scans are deterministic and keep working offline /
//! on-prem), and maps the catalog into the corpus controls the mapping engine
//! consumes. The producer owns the catalog; we own the assessment — this is the
//! ingest half of the loop.

use std::path::PathBuf;

use secrecy::{ExposeSecret, SecretString};

use compliance_core::error::CoreError;
use compliance_core::models::onboarding::ComplianceFramework;
use compliance_core::models::oscal::OscalDocument;
use compliance_core::traits::{Control, ControlQuery, ControlsProvider};

/// A [`ControlsProvider`] that pulls the OSCAL catalog from breakpilot-compliance
/// and snapshots it locally for deterministic / offline reuse.
pub struct OscalControlsProvider {
    http: reqwest::Client,
    base_url: String,
    token: Option<SecretString>,
    snapshot_dir: PathBuf,
}

impl OscalControlsProvider {
    /// Create a provider. `base_url` is the breakpilot-compliance root (e.g.
    /// `http://backend-compliance:8002`); `snapshot_dir` is where catalog
    /// snapshots are written so a later scan can reuse them without the network.
    pub fn new(
        http: reqwest::Client,
        base_url: impl Into<String>,
        token: Option<SecretString>,
        snapshot_dir: impl Into<PathBuf>,
    ) -> Self {
        Self {
            http,
            base_url: base_url.into(),
            token,
            snapshot_dir: snapshot_dir.into(),
        }
    }

    fn catalog_url(&self, framework: &str) -> String {
        format!(
            "{}/api/compliance/v1/oscal/catalog?framework={framework}",
            self.base_url.trim_end_matches('/')
        )
    }

    fn snapshot_path(&self, framework: &str) -> PathBuf {
        self.snapshot_dir
            .join(format!("oscal-catalog-{framework}.json"))
    }

    /// Fetch the raw catalog bytes for a framework token over HTTP.
    async fn fetch_raw(&self, framework: &str) -> Result<Vec<u8>, CoreError> {
        let mut req = self.http.get(self.catalog_url(framework));
        if let Some(token) = &self.token {
            req = req.bearer_auth(token.expose_secret());
        }
        let resp = req
            .send()
            .await
            .map_err(|e| CoreError::Http(e.to_string()))?;
        if !resp.status().is_success() {
            return Err(CoreError::Http(format!(
                "catalog fetch for {framework} returned HTTP {}",
                resp.status()
            )));
        }
        resp.bytes()
            .await
            .map(|b| b.to_vec())
            .map_err(|e| CoreError::Http(e.to_string()))
    }

    /// Write a catalog snapshot atomically (temp file + rename).
    async fn write_snapshot(&self, framework: &str, raw: &[u8]) -> Result<(), CoreError> {
        tokio::fs::create_dir_all(&self.snapshot_dir).await?;
        let path = self.snapshot_path(framework);
        let tmp = path.with_extension("json.tmp");
        tokio::fs::write(&tmp, raw).await?;
        tokio::fs::rename(&tmp, &path).await?;
        Ok(())
    }

    /// Read a previously written snapshot, if one exists.
    async fn read_snapshot(&self, framework: &str) -> Result<Option<OscalDocument>, CoreError> {
        match tokio::fs::read(self.snapshot_path(framework)).await {
            Ok(raw) => Ok(Some(serde_json::from_slice(&raw)?)),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    /// Load the catalog for a framework token: fetch fresh + snapshot the exact
    /// bytes; on network failure, fall back to the last snapshot so scans run.
    async fn load_token(&self, framework: &str) -> Result<OscalDocument, CoreError> {
        match self.fetch_raw(framework).await {
            Ok(raw) => {
                let doc: OscalDocument = serde_json::from_slice(&raw)?;
                if let Err(e) = self.write_snapshot(framework, &raw).await {
                    tracing::warn!(framework, error = %e, "failed to write OSCAL snapshot");
                }
                Ok(doc)
            }
            Err(fetch_err) => match self.read_snapshot(framework).await? {
                Some(doc) => {
                    tracing::warn!(
                        framework, error = %fetch_err,
                        "OSCAL catalog fetch failed; falling back to snapshot"
                    );
                    Ok(doc)
                }
                None => Err(fetch_err),
            },
        }
    }

    /// Load the OSCAL catalog for a compliance framework.
    pub async fn load(&self, framework: ComplianceFramework) -> Result<OscalDocument, CoreError> {
        self.load_token(&framework.to_string()).await
    }

    /// Load the code-checkable master-controls catalog
    /// (`?framework=master-controls`).
    pub async fn load_master_controls(&self) -> Result<OscalDocument, CoreError> {
        self.load_token("master-controls").await
    }
}

/// Order controls whose title/text mention the query context first (stable), then
/// truncate to the requested limit. Naive relevance — refined when the assessment
/// layer lands.
fn rank_and_truncate(mut controls: Vec<Control>, context: &str, limit: usize) -> Vec<Control> {
    if !context.is_empty() {
        let needle = context.to_lowercase();
        controls.sort_by_key(|c| {
            let hit =
                c.title.to_lowercase().contains(&needle) || c.text.to_lowercase().contains(&needle);
            u8::from(!hit)
        });
    }
    controls.truncate(limit);
    controls
}

impl ControlsProvider for OscalControlsProvider {
    fn name(&self) -> &str {
        "breakpilot-oscal"
    }

    async fn controls(&self, query: &ControlQuery<'_>) -> Result<Vec<Control>, CoreError> {
        let mut out: Vec<Control> = Vec::new();
        for &framework in query.frameworks {
            match self.load(framework).await {
                Ok(doc) => out.extend(doc.to_controls()),
                Err(e) => {
                    tracing::warn!(%framework, error = %e, "skipping framework: catalog unavailable")
                }
            }
        }
        Ok(rank_and_truncate(out, query.context, query.limit))
    }
}

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

    const MINI_CATALOG: &str = r#"{"catalog":{"uuid":"u","metadata":{"title":"T",
      "version":"1.0.0","oscal-version":"1.1.2","props":[{"name":"framework","value":"cra"}]},
      "groups":[{"id":"g","title":"G","controls":[{"id":"cra-ai-1","title":"MFA",
      "props":[],"parts":[{"name":"statement","prose":"require mfa"}]}]}]}}"#;

    fn provider(dir: &std::path::Path) -> OscalControlsProvider {
        OscalControlsProvider::new(reqwest::Client::new(), "http://unused/", None, dir)
    }

    #[test]
    fn builds_catalog_url_and_snapshot_path() {
        let p = provider(std::path::Path::new("/snap"));
        assert_eq!(
            p.catalog_url("cra"),
            "http://unused/api/compliance/v1/oscal/catalog?framework=cra"
        );
        assert_eq!(
            p.snapshot_path("cra"),
            std::path::Path::new("/snap/oscal-catalog-cra.json")
        );
    }

    #[test]
    fn ranks_context_hits_first_then_truncates() {
        let mk = |id: &str, title: &str| Control {
            id: id.into(),
            framework: ComplianceFramework::Cra,
            title: title.into(),
            text: String::new(),
            source: None,
        };
        let controls = vec![
            mk("a", "logging policy"),
            mk("b", "multi-factor auth"),
            mk("c", "backup"),
        ];
        let ranked = rank_and_truncate(controls, "auth", 2);
        assert_eq!(ranked.len(), 2);
        assert_eq!(ranked[0].id, "b"); // the "auth" hit floats to the top
    }

    #[tokio::test]
    async fn snapshot_round_trip_and_offline_fallback() {
        let dir = std::env::temp_dir().join(format!("oscal-test-{}", uuid::Uuid::new_v4()));
        let p = provider(&dir);
        assert!(p.read_snapshot("cra").await.unwrap().is_none());
        p.write_snapshot("cra", MINI_CATALOG.as_bytes())
            .await
            .unwrap();
        let doc = p.read_snapshot("cra").await.unwrap().unwrap();
        assert_eq!(doc.to_controls().len(), 1);
        assert_eq!(doc.framework(), Some(ComplianceFramework::Cra));
        let _ = std::fs::remove_dir_all(&dir);
    }
}