compliance-scanner-agent/compliance-agent/src/classify/language.rs

//! Heuristic target-type classification from artifact kinds and source markers.
//!
//! Complements the tramiton firmware detector: this handles web / backend /
//! mobile / desktop / PLC by sniffing manifest files and file extensions in the
//! ingested code trees, plus strong priors from the artifact kinds themselves
//! (a PLC-project artifact is a PLC target; an `.ipa` is an iOS app).

use std::collections::HashSet;
use std::fs;
use std::path::Path;

use compliance_core::error::CoreError;
use compliance_core::models::{ArtifactKind, DetectedFact, TargetType};
use compliance_core::traits::{ClassificationInput, ClassifierVerdict, TargetClassifier};

/// Max directory depth scanned for marker files.
const SCAN_DEPTH: usize = 2;

/// Markers collected from a code tree.
#[derive(Default)]
struct Markers {
    files: HashSet<String>,
    dirs: HashSet<String>,
    exts: HashSet<String>,
}

impl Markers {
    fn has_file(&self, name: &str) -> bool {
        self.files.contains(name)
    }
    fn has_ext(&self, ext: &str) -> bool {
        self.exts.contains(ext)
    }
    fn any_dir_ends_with(&self, suffix: &str) -> bool {
        self.dirs.iter().any(|d| d.ends_with(suffix))
    }
}

/// Recursively collect marker file/dir/extension names up to [`SCAN_DEPTH`].
fn collect_markers(root: &Path) -> Markers {
    let mut m = Markers::default();
    scan_dir(root, 0, &mut m);
    m
}

fn scan_dir(dir: &Path, depth: usize, m: &mut Markers) {
    let Ok(entries) = fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        let name = entry.file_name().to_string_lossy().to_lowercase();
        if path.is_dir() {
            m.dirs.insert(name);
            if depth < SCAN_DEPTH {
                scan_dir(&path, depth + 1, m);
            }
        } else {
            if let Some(ext) = path.extension() {
                m.exts.insert(ext.to_string_lossy().to_lowercase());
            }
            m.files.insert(name);
        }
    }
}

/// Whether a `package.json` at `root` looks like a front-end app.
fn package_json_is_frontend(root: &Path) -> bool {
    let Ok(content) = fs::read_to_string(root.join("package.json")) else {
        return false;
    };
    let c = content.to_lowercase();
    ["react", "next", "vue", "@angular", "svelte", "vite"]
        .iter()
        .any(|f| c.contains(f))
}

/// The heuristic classifier: artifact-kind priors + source-tree markers.
pub struct HeuristicClassifier;

impl HeuristicClassifier {
    /// Verdicts from the artifact kinds alone (no filesystem needed).
    fn kind_priors(&self, input: &ClassificationInput<'_>) -> Vec<ClassifierVerdict> {
        let mut out = Vec::new();
        for a in input.artifacts {
            let lower = a.source_ref.to_lowercase();
            match a.kind {
                ArtifactKind::PlcProject => out.push(verdict(
                    TargetType::PlcSps,
                    0.85,
                    "PLC project artifact",
                    vec![],
                )),
                ArtifactKind::MobilePackage => {
                    let (tt, why) = if lower.ends_with(".ipa") {
                        (TargetType::IosApp, "iOS package (.ipa)")
                    } else {
                        (TargetType::AndroidApp, "Android package (.apk/.aab)")
                    };
                    out.push(verdict(tt, 0.85, why, vec![]));
                }
                ArtifactKind::ContainerImage => out.push(verdict(
                    TargetType::BackendService,
                    0.4,
                    "container image",
                    vec![],
                )),
                ArtifactKind::FirmwareImage => out.push(verdict(
                    TargetType::FirmwareBareMetal,
                    0.35,
                    "firmware image (pending tramiton detection)",
                    vec![],
                )),
                ArtifactKind::LiveUrl if input.artifacts.len() == 1 => {
                    out.push(verdict(TargetType::WebApp, 0.3, "live URL only", vec![]))
                }
                _ => {}
            }
        }
        out
    }

    /// Verdicts from scanning the ingested code trees for manifest markers.
    fn source_verdicts(&self, input: &ClassificationInput<'_>) -> Vec<ClassifierVerdict> {
        let mut out = Vec::new();
        for a in input.artifacts {
            if !matches!(a.kind, ArtifactKind::GitRepo | ArtifactKind::SourceArchive) {
                continue;
            }
            let Some(path) = input.working_paths.get(&a.id) else {
                continue;
            };
            let m = collect_markers(path);

            // Mobile (checked first — strongest signal).
            if m.has_file("androidmanifest.xml") || m.has_ext("apk") || m.has_ext("aab") {
                out.push(verdict(
                    TargetType::AndroidApp,
                    0.8,
                    "Android manifest / gradle",
                    facts_lang("kotlin/java"),
                ));
            }
            if m.any_dir_ends_with(".xcodeproj")
                || m.has_file("info.plist")
                || m.has_file("podfile")
                || m.has_ext("ipa")
            {
                out.push(verdict(
                    TargetType::IosApp,
                    0.8,
                    "Xcode project / Info.plist",
                    facts_lang("swift/objc"),
                ));
            }
            // Desktop.
            if m.has_ext("sln")
                || m.has_ext("csproj")
                || m.has_ext("vcxproj")
                || m.has_ext("desktop")
            {
                out.push(verdict(
                    TargetType::DesktopApp,
                    0.7,
                    "desktop project files",
                    facts_lang("dotnet/native"),
                ));
            }
            // PLC.
            if m.has_ext("st") {
                out.push(verdict(
                    TargetType::PlcSps,
                    0.8,
                    "Structured Text sources",
                    facts_lang("iec-61131-3"),
                ));
            }
            // Web vs backend from package.json.
            if m.has_file("package.json") {
                if package_json_is_frontend(path) {
                    out.push(verdict(
                        TargetType::WebApp,
                        0.65,
                        "package.json with a front-end framework",
                        facts_lang("javascript"),
                    ));
                } else {
                    out.push(verdict(
                        TargetType::BackendService,
                        0.55,
                        "package.json (no front-end framework)",
                        facts_lang("javascript"),
                    ));
                }
            }
            // Backend languages.
            for (file, lang) in [
                ("cargo.toml", "rust"),
                ("go.mod", "go"),
                ("pom.xml", "java"),
                ("requirements.txt", "python"),
                ("pyproject.toml", "python"),
            ] {
                if m.has_file(file) {
                    out.push(verdict(
                        TargetType::BackendService,
                        0.6,
                        "backend build manifest",
                        facts_lang(lang),
                    ));
                }
            }
            // Container-only.
            if m.has_file("dockerfile") && out.is_empty() {
                out.push(verdict(
                    TargetType::BackendService,
                    0.4,
                    "Dockerfile",
                    facts_lang("container"),
                ));
            }
        }
        out
    }
}

impl TargetClassifier for HeuristicClassifier {
    fn name(&self) -> &str {
        "heuristic"
    }

    async fn classify(
        &self,
        input: &ClassificationInput<'_>,
    ) -> Result<Vec<ClassifierVerdict>, CoreError> {
        let mut out = self.kind_priors(input);
        out.extend(self.source_verdicts(input));
        Ok(out)
    }
}

fn verdict(
    target_type: TargetType,
    confidence: f32,
    rationale: &str,
    facts: Vec<DetectedFact>,
) -> ClassifierVerdict {
    ClassifierVerdict {
        target_type,
        confidence,
        facts,
        rationale: rationale.to_string(),
    }
}

fn facts_lang(lang: &str) -> Vec<DetectedFact> {
    vec![DetectedFact::new("language", lang, "heuristic")]
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;
    use compliance_core::models::Artifact;
    use std::collections::HashMap;
    use std::path::PathBuf;

    struct Scratch(PathBuf);
    impl Scratch {
        fn new() -> Self {
            let p = std::env::temp_dir().join(format!("cs-classify-{}", uuid::Uuid::new_v4()));
            fs::create_dir_all(&p).expect("mkdir");
            Self(p)
        }
    }
    impl Drop for Scratch {
        fn drop(&mut self) {
            let _ = fs::remove_dir_all(&self.0);
        }
    }

    async fn classify_tree(setup: impl FnOnce(&Path)) -> Vec<ClassifierVerdict> {
        let scratch = Scratch::new();
        setup(&scratch.0);
        let artifact = Artifact::git_repo("https://git/x", "main");
        let mut wp = HashMap::new();
        wp.insert(artifact.id.clone(), scratch.0.clone());
        let artifacts = vec![artifact];
        let input = ClassificationInput {
            artifacts: &artifacts,
            working_paths: &wp,
            description: None,
        };
        HeuristicClassifier
            .classify(&input)
            .await
            .expect("classify")
    }

    #[tokio::test]
    async fn frontend_package_json_is_webapp() {
        let v = classify_tree(|root| {
            fs::write(
                root.join("package.json"),
                r#"{"dependencies":{"react":"18"}}"#,
            )
            .unwrap();
        })
        .await;
        assert!(v.iter().any(|x| x.target_type == TargetType::WebApp));
    }

    #[tokio::test]
    async fn cargo_toml_is_backend() {
        let v = classify_tree(|root| {
            fs::write(root.join("Cargo.toml"), "[package]\nname='x'").unwrap();
        })
        .await;
        assert!(v
            .iter()
            .any(|x| x.target_type == TargetType::BackendService));
    }

    #[tokio::test]
    async fn android_manifest_is_android() {
        let v = classify_tree(|root| {
            fs::write(root.join("AndroidManifest.xml"), "<manifest/>").unwrap();
        })
        .await;
        assert!(v.iter().any(|x| x.target_type == TargetType::AndroidApp));
    }

    #[tokio::test]
    async fn structured_text_is_plc() {
        let v = classify_tree(|root| {
            fs::write(root.join("main.st"), "PROGRAM main END_PROGRAM").unwrap();
        })
        .await;
        assert!(v.iter().any(|x| x.target_type == TargetType::PlcSps));
    }

    #[tokio::test]
    async fn ipa_artifact_prior_is_ios() {
        let artifacts = vec![Artifact::mobile_package("app.ipa")];
        let wp = HashMap::new();
        let input = ClassificationInput {
            artifacts: &artifacts,
            working_paths: &wp,
            description: None,
        };
        let v = HeuristicClassifier
            .classify(&input)
            .await
            .expect("classify");
        assert!(v.iter().any(|x| x.target_type == TargetType::IosApp));
    }
}