agentsec-core 0.2.0

//! BlackList check: classify installed MCP server names against a
//! [`crate::registry::Registry`] of known-good entries.
//!
//! This is a **demo-level** signature check: it reads MCP server names
//! from `.mcp.json` (cwd) and `.claude.json` (under [`crate::Paths::user_home`])
//! at call time, then assigns each name one of three verdicts:
//!
//! - [`UnknownVerdictKind::KnownGood`] — exact name match in the registry.
//! - [`UnknownVerdictKind::LikelyTyposquat`] — Levenshtein distance == 1 to a
//!   registry entry (not identical).  High-signal, warn severity.
//! - [`UnknownVerdictKind::InformationalTyposquat`] — Levenshtein distance == 2
//!   to a registry entry (not identical, and name length > 3).  Lower-signal,
//!   info severity.
//! - [`UnknownVerdictKind::Unknown`] — no match; the user has installed
//!   an MCP server AgentSec doesn't recognise. This is **neutral**, not a
//!   block — it just surfaces "we've never seen this".
//!
//! The classify step does not consume the inventory snapshot — it
//! re-reads the source JSON because the snapshot only stores hashes, not
//! the full structure.
//!
//! [`classify`] groups rows by `name`, deduplicates across files, and
//! returns one [`UnknownVerdict`] per distinct server name with `count`
//! (number of config files it appeared in) and `paths` (sorted list of
//! those files).

use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::PathBuf;

use serde::{Deserialize, Serialize};

use crate::Paths;
use crate::registry::Registry;

const TYPOSQUAT_DISTANCE_MAX: usize = 2;

/// One row of [`classify`]'s aggregated output.
///
/// A single `name` may appear in multiple config files.  `count` is the
/// number of distinct files and `paths` lists them in sorted order.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct UnknownVerdict {
    /// MCP server name (the key under `mcpServers.<name>`).
    pub name: String,
    /// Classification of `name` against the registry.
    pub verdict: UnknownVerdictKind,
    /// Free-text rationale, suitable for one-line surface in the Markdown
    /// report. Echoes the registry [`crate::registry::RegistryEntry::source`]
    /// when relevant.
    pub reason: String,
    /// Number of config files this name was found in.
    pub count: usize,
    /// Sorted list of config files this name was found in.
    pub paths: Vec<String>,
}

/// Classification levels.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum UnknownVerdictKind {
    /// Exact match in the registry.
    KnownGood,
    /// Levenshtein distance == 1 to a registry entry (warn severity).
    LikelyTyposquat { closest: String, distance: usize },
    /// Levenshtein distance == 2 to a registry entry (info severity).
    /// Suppressed for names with ≤ 3 characters.
    InformationalTyposquat { closest: String, distance: usize },
    /// No match at any distance ≤ 2 (or suppressed short-name case).
    Unknown,
}

/// Read MCP server names from `.mcp.json` and `.claude.json`, classify each
/// one against `registry`, and return one aggregated row per distinct name.
///
/// File-read failures are silently skipped — a missing `.mcp.json` is the
/// common case, and a malformed `.claude.json` shouldn't fail the whole
/// classification.
///
/// Output is sorted by `name` for stable reporting.
pub fn classify(paths: &Paths, registry: &Registry) -> Vec<UnknownVerdict> {
    // Collect (path, name) pairs — BTreeSet gives dedup + sorted order.
    let mut named: BTreeSet<(String, String)> = BTreeSet::new();

    // ── project-rooted `.mcp.json` (cwd-relative) ─────────────────────
    let project_mcp = PathBuf::from(".mcp.json");
    extract_names(&project_mcp, &mut named, /*per_project=*/ false);

    // ── home-rooted `.claude.json` (top-level + per-project blocks) ───
    let home_config = paths.user_home.join(".claude.json");
    extract_names(&home_config, &mut named, /*per_project=*/ true);

    // Group by name: BTreeMap keeps names sorted.
    let mut by_name: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
    for (path, name) in named {
        by_name.entry(name).or_default().insert(path);
    }

    by_name
        .into_iter()
        .map(|(name, path_set)| {
            let paths_vec: Vec<String> = path_set.into_iter().collect();
            let count = paths_vec.len();
            // Use the first path as the representative for classification;
            // the verdict depends only on the name + registry.
            let rep = paths_vec.first().map_or("", String::as_str);
            let mut v = classify_one(rep, &name, registry);
            v.count = count;
            v.paths = paths_vec;
            v
        })
        .collect()
}

fn classify_one(path: &str, name: &str, registry: &Registry) -> UnknownVerdict {
    if let Some(entry) = registry.get(name) {
        return UnknownVerdict {
            name: name.to_string(),
            verdict: UnknownVerdictKind::KnownGood,
            reason: format!("registered (source: {})", entry.source),
            count: 1,
            paths: vec![path.to_string()],
        };
    }
    if let Some((entry, distance)) = registry.closest(name, TYPOSQUAT_DISTANCE_MAX) {
        // Suppress distance-2 hits for very short names (≤ 3 chars) to avoid
        // false positives like "git" → "filesystem" (not a real threat model).
        let char_count = name.chars().count();
        if distance == 2 && char_count <= 3 {
            // Fall through to Unknown below.
        } else {
            let verdict = if distance == 1 {
                UnknownVerdictKind::LikelyTyposquat {
                    closest: entry.name.clone(),
                    distance,
                }
            } else {
                UnknownVerdictKind::InformationalTyposquat {
                    closest: entry.name.clone(),
                    distance,
                }
            };
            return UnknownVerdict {
                name: name.to_string(),
                verdict,
                reason: format!("close to `{}` (distance {})", entry.name, distance),
                count: 1,
                paths: vec![path.to_string()],
            };
        }
    }
    UnknownVerdict {
        name: name.to_string(),
        verdict: UnknownVerdictKind::Unknown,
        reason: "not in registry".to_string(),
        count: 1,
        paths: vec![path.to_string()],
    }
}

/// Read `path` as JSON and extract MCP server names from top-level
/// `mcpServers` (and, if `per_project`, from each
/// `projects.<p>.mcpServers` block too). Adds `(display-path, name)`
/// tuples to `out`. Silently no-ops on missing / unparseable files.
fn extract_names(path: &PathBuf, out: &mut BTreeSet<(String, String)>, per_project: bool) {
    let Ok(body) = fs::read_to_string(path) else {
        return;
    };
    let Ok(json) = serde_json::from_str::<serde_json::Value>(&body) else {
        return;
    };
    let display = path.display().to_string();
    if let Some(top) = json
        .get("mcpServers")
        .and_then(serde_json::Value::as_object)
    {
        for name in top.keys() {
            out.insert((display.clone(), name.clone()));
        }
    }
    if per_project {
        if let Some(projects) = json.get("projects").and_then(serde_json::Value::as_object) {
            for (proj_name, proj_val) in projects {
                if let Some(servers) = proj_val
                    .get("mcpServers")
                    .and_then(serde_json::Value::as_object)
                {
                    for name in servers.keys() {
                        let scoped = format!("{display}#projects.{proj_name}.mcpServers");
                        out.insert((scoped, name.clone()));
                    }
                }
            }
        }
    }
}

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

    fn write_claude_json(tmp: &tempfile::TempDir, body: &str) {
        std::fs::write(tmp.path().join(".claude.json"), body).unwrap();
    }

    fn paths_for(tmp: &tempfile::TempDir) -> Paths {
        Paths {
            home: tmp.path().to_path_buf(),
            user_home: tmp.path().to_path_buf(),
        }
    }

    fn small_registry() -> Registry {
        Registry::from_entries(vec![
            RegistryEntry {
                name: "filesystem".into(),
                source: "test-source".into(),
            },
            RegistryEntry {
                name: "github".into(),
                source: "test-source".into(),
            },
        ])
    }

    #[test]
    fn known_good_is_recognised() {
        let tmp = tempfile::tempdir().unwrap();
        write_claude_json(&tmp, r#"{"mcpServers":{"filesystem":{"command":"x"}}}"#);
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert_eq!(verdicts.len(), 1);
        assert_eq!(verdicts[0].name, "filesystem");
        assert_eq!(verdicts[0].verdict, UnknownVerdictKind::KnownGood);
    }

    #[test]
    fn typosquat_is_flagged() {
        let tmp = tempfile::tempdir().unwrap();
        // "filesystme" — transposition of last two chars from "filesystem",
        // Levenshtein distance 2 → InformationalTyposquat.
        write_claude_json(&tmp, r#"{"mcpServers":{"filesystme":{"command":"x"}}}"#);
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert_eq!(verdicts.len(), 1);
        match &verdicts[0].verdict {
            UnknownVerdictKind::InformationalTyposquat { closest, distance } => {
                assert_eq!(closest, "filesystem");
                assert_eq!(*distance, 2);
            }
            other => panic!("expected InformationalTyposquat, got {other:?}"),
        }
    }

    #[test]
    fn unknown_name_is_unknown() {
        let tmp = tempfile::tempdir().unwrap();
        write_claude_json(
            &tmp,
            r#"{"mcpServers":{"my-private-tool":{"command":"x"}}}"#,
        );
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert_eq!(verdicts.len(), 1);
        assert_eq!(verdicts[0].verdict, UnknownVerdictKind::Unknown);
    }

    #[test]
    fn per_project_servers_are_classified() {
        let tmp = tempfile::tempdir().unwrap();
        write_claude_json(
            &tmp,
            r#"{
                "mcpServers": {"github": {"command":"x"}},
                "projects": {
                    "/some/proj": {
                        "mcpServers": {"my-tool": {"command":"y"}}
                    }
                }
            }"#,
        );
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert_eq!(verdicts.len(), 2);
        let names: Vec<&str> = verdicts.iter().map(|v| v.name.as_str()).collect();
        assert!(names.contains(&"github"));
        assert!(names.contains(&"my-tool"));
    }

    #[test]
    fn missing_claude_json_yields_empty_result() {
        let tmp = tempfile::tempdir().unwrap();
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert!(verdicts.is_empty());
    }

    #[test]
    fn malformed_json_silently_skipped() {
        let tmp = tempfile::tempdir().unwrap();
        write_claude_json(&tmp, "not even close to JSON");
        let verdicts = classify(&paths_for(&tmp), &small_registry());
        assert!(verdicts.is_empty());
    }

    // ── A-2 heuristic tests ───────────────────────────────────────────────

    #[test]
    fn classify_one_short_name_distance_2_suppressed() {
        // "git" has length 3, closest registry entry is "github" (distance 3)
        // or "filesystem" (distance > 2). Use a custom registry with a
        // 2-distance match to exercise the short-name suppression.
        let reg = Registry::from_entries(vec![RegistryEntry {
            name: "liv".into(),
            source: "test".into(),
        }]);
        // "git" → "liv": levenshtein("git","liv") = 3 → actually no match at max 2.
        // Use "gi" (2 chars) and "github" (distance 5) — too far.
        // Directly test classify_one: name "gxt", registry has "git", distance 2.
        let reg2 = Registry::from_entries(vec![RegistryEntry {
            name: "git".into(),
            source: "test".into(),
        }]);
        // "gxt" vs "git" = distance 1 (transposition: x→i) — actually 1, not 2.
        // Use "ab" vs "cd": distance 2, both len 2 → suppressed.
        let reg3 = Registry::from_entries(vec![RegistryEntry {
            name: "ab".into(),
            source: "test".into(),
        }]);
        let v = classify_one("path", "cd", &reg3);
        // "cd" len=2 ≤ 3, distance 2 → suppressed to Unknown
        assert_eq!(v.verdict, UnknownVerdictKind::Unknown);

        // Sanity: if registry has "ab", "abc" (len 3 ≤ 3) close at distance 2
        // to "axy" → suppressed
        let reg4 = Registry::from_entries(vec![RegistryEntry {
            name: "abc".into(),
            source: "test".into(),
        }]);
        // "axz" vs "abc" = 2 subs → distance 2, len("axz")=3 ≤ 3 → suppressed
        let v2 = classify_one("path", "axz", &reg4);
        assert_eq!(v2.verdict, UnknownVerdictKind::Unknown);

        // Suppress unused-var warning for reg, reg2 by using them.
        let _ = reg;
        let _ = reg2;
    }

    #[test]
    fn classify_severity_split() {
        let reg = Registry::from_entries(vec![RegistryEntry {
            name: "filesystem".into(),
            source: "test".into(),
        }]);
        // distance 1 → LikelyTyposquat
        let v1 = classify_one("p", "filesytem", &reg); // filesytem vs filesystem: 1 missing char
        match &v1.verdict {
            UnknownVerdictKind::LikelyTyposquat { distance, .. } => assert_eq!(*distance, 1),
            other => panic!("expected LikelyTyposquat, got {other:?}"),
        }
        // distance 2, long name → InformationalTyposquat
        let v2 = classify_one("p", "filesystXY", &reg); // 2 substitutions
        match &v2.verdict {
            UnknownVerdictKind::InformationalTyposquat { distance, .. } => {
                assert_eq!(*distance, 2);
            }
            other => panic!("expected InformationalTyposquat, got {other:?}"),
        }
    }

    #[test]
    fn classify_aggregates_dedup_count() {
        use std::fs;
        let tmp = tempfile::tempdir().unwrap();
        let paths = paths_for(&tmp);
        // Write .claude.json with "my-tool" in top-level and in two projects.
        write_claude_json(
            &tmp,
            r#"{
                "mcpServers": {"my-tool": {"command":"x"}},
                "projects": {
                    "/proj1": {"mcpServers": {"my-tool": {"command":"y"}}},
                    "/proj2": {"mcpServers": {"my-tool": {"command":"z"}}}
                }
            }"#,
        );
        let reg = Registry::from_entries(vec![]);
        let verdicts = classify(&paths, &reg);
        // All occurrences of "my-tool" are aggregated into one row.
        let v = verdicts.iter().find(|v| v.name == "my-tool").unwrap();
        // 3 occurrences: top-level + 2 project scopes
        assert_eq!(v.count, 3);
        assert_eq!(v.paths.len(), 3);
        // paths are sorted
        assert!(v.paths.windows(2).all(|w| w[0] <= w[1]));

        // Suppress unused import warning
        let _ = fs::metadata(tmp.path());
    }
}