//! Capability-tier toxic-flow analysis (, Snyk W015-W020 / AgentSeal
//! Toxic Flow Mapping).
//!
//! Why this is a separate module: the surface checks in [`crate::security::surface`]
//! catch the directive half of a poisoning attack (SEC-003 tells the model to
//! exfiltrate, SEC-009 flags an unannotated destructive tool). They do not score
//! the latent capability of a catalog before any payload fires. A toxic flow is
//! structural: a tool that reads untrusted content, plus a tool that can
//! exfiltrate sensitive data or destroy state, is a complete exfiltration
//! chain waiting for a prompt-injected instruction to wire the two together.
//! That risk exists at classification time, not at call time.
//!
//! The pairing is relational (it needs two tools at once), so like the
//! namespace and integrity families this lives outside the
//! per-[`Definition`](crate::security::Definition) engine as a standalone function over
//! the full snapshot.
//!
//! This is a conservative keyword heuristic, so detectability is partial: a
//! tool named or described in terms that dodge the keyword lists below will not
//! be classified, and a benign tool whose name happens to contain a keyword may
//! be over-classified. The lists favor recall on the obvious cases over
//! precision on the edge cases, because a missed toxic flow is the worse error.

use serde_json::Value;

use super::{CheckCategory, DefinitionKind, SecurityFinding, SecuritySeverity};

/// A latent capability tier a tool can hold, scored before any call fires.
///
/// A tool can hold zero or more tiers at once (a `read_repo_file` tool is a
/// workspace reader; a `fetch_and_delete` tool is both an untrusted-content
/// source and a destructive sink). The tiers split exposure (what a tool can
/// read out) from destruction (what a tool can write or run), and split the
/// blast radius (sensitive vs workspace, remote vs local) so the pairing rule
/// can weigh a credential reader more heavily than a file reader.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum CapabilityTier {
    /// Reads sensitive data: email, credentials, vault entries, financial or
    /// secret material. The highest-value thing an attacker wants exfiltrated.
    SensitiveDataExposure,
    /// Reads workspace data: files, paths, repositories, source code. Lower
    /// value than secrets but still an exfiltration target.
    WorkspaceDataExposure,
    /// Can destroy or move state remotely: delete, drop, execute a command,
    /// transfer funds, make a payment. The remote-blast destructive sink.
    Destructive,
    /// Can destroy local state: remove or unlink a file, format a disk, write
    /// to a local path. The local-blast destructive sink.
    LocalDestructive,
    /// Pulls untrusted content into the context: fetches a URL, reads the web,
    /// loads third-party or external content. The injection entry point.
    UntrustedContentSource,
}

impl CapabilityTier {
    /// Lowercase label used in finding evidence and the catalog.
    pub fn as_str(self) -> &'static str {
        match self {
            CapabilityTier::SensitiveDataExposure => "sensitive-data-exposure",
            CapabilityTier::WorkspaceDataExposure => "workspace-data-exposure",
            CapabilityTier::Destructive => "destructive",
            CapabilityTier::LocalDestructive => "local-destructive",
            CapabilityTier::UntrustedContentSource => "untrusted-content-source",
        }
    }

    /// True when this tier is a sink an attacker can route stolen data or
    /// destruction through. The pairing rule treats these tiers as the
    /// dangerous half of a toxic flow; reaching one of them after an untrusted
    /// read completes the chain.
    fn is_sink(self) -> bool {
        matches!(
            self,
            CapabilityTier::SensitiveDataExposure
                | CapabilityTier::Destructive
                | CapabilityTier::LocalDestructive
        )
    }
}

// ---------------------------------------------------------------------------
// Keyword heuristic lists
// ---------------------------------------------------------------------------
//
// Each list is the conservative keyword set for one tier. A keyword matches as
// a lowercase substring of the tool's name, description, and input-schema
// property names combined. Substring matching is deliberate so `read_file`
// matches `file` and `delete_record` matches `delete`; the cost is the rare
// false positive (a tool about a "filename" matching `file`), which is why the
// per-tool tier finding is informational rather than a defect.

/// Sensitive-data-exposure keywords: things that read high-value secrets.
const SENSITIVE_KEYWORDS: &[&str] = &[
    "email",
    "credential",
    "vault",
    "financial",
    "secret",
    "password",
    "token",
    "api key",
    "api_key",
    "apikey",
    "ssh",
    "private key",
    "private_key",
    "wallet",
    "bank",
    "ssn",
    "pii",
];

/// Workspace-data-exposure keywords: things that read project or local data.
const WORKSPACE_KEYWORDS: &[&str] = &[
    "file",
    "path",
    "repo",
    "repository",
    "code",
    "directory",
    "folder",
    "document",
    "workspace",
    "source",
];

/// Destructive keywords: things that destroy or move state remotely.
const DESTRUCTIVE_KEYWORDS: &[&str] = &[
    "delete",
    "drop",
    "exec",
    "command",
    "transfer",
    "pay",
    "purge",
    "destroy",
    "wipe",
    "truncate",
    "revoke",
    "send money",
    "withdraw",
];

/// Local-destructive keywords: things that destroy local state.
const LOCAL_DESTRUCTIVE_KEYWORDS: &[&str] = &[
    "rm ",
    "rm -",
    "unlink",
    "format",
    "rmdir",
    "shred",
    "local path",
    "local file",
    "overwrite local",
];

/// Untrusted-content-source keywords: things that pull external content in.
const UNTRUSTED_SOURCE_KEYWORDS: &[&str] = &[
    "fetch",
    "http",
    "url",
    "web",
    "crawl",
    "scrape",
    "download",
    "browse",
    "read external",
    "third-party",
    "third party",
    "remote content",
    "rss",
    "webhook",
];

/// Classify one tool object into the capability tiers its name, description,
/// and input-schema property names imply.
///
/// `tool` is a single entry from a snapshot's `tools` array. The function
/// lowercases and concatenates the searchable text once, then tests each tier's
/// keyword list against it. Tiers come back sorted and deduplicated, so a tool
/// that matches the same tier through two keywords still reports it once. A
/// non-tool value or a tool with no searchable text returns an empty vector.
///
/// ```
/// use mcptest_core::security::{classify_tool, CapabilityTier};
/// use serde_json::json;
///
/// let tool = json!({ "name": "fetch_url", "description": "Fetch a web page." });
/// let tiers = classify_tool(&tool);
/// assert!(tiers.contains(&CapabilityTier::UntrustedContentSource));
/// ```
pub fn classify_tool(tool: &Value) -> Vec<CapabilityTier> {
    let haystack = searchable_text(tool);
    if haystack.is_empty() {
        return Vec::new();
    }
    let mut tiers = Vec::new();
    if matches_any(&haystack, SENSITIVE_KEYWORDS) {
        tiers.push(CapabilityTier::SensitiveDataExposure);
    }
    if matches_any(&haystack, WORKSPACE_KEYWORDS) {
        tiers.push(CapabilityTier::WorkspaceDataExposure);
    }
    if matches_any(&haystack, DESTRUCTIVE_KEYWORDS) {
        tiers.push(CapabilityTier::Destructive);
    }
    if matches_any(&haystack, LOCAL_DESTRUCTIVE_KEYWORDS) {
        tiers.push(CapabilityTier::LocalDestructive);
    }
    if matches_any(&haystack, UNTRUSTED_SOURCE_KEYWORDS) {
        tiers.push(CapabilityTier::UntrustedContentSource);
    }
    tiers.sort();
    tiers.dedup();
    tiers
}

/// Lowercased concatenation of a tool's name, description, and input-schema
/// property names, the text the keyword lists search.
fn searchable_text(tool: &Value) -> String {
    let mut parts = Vec::new();
    if let Some(name) = tool.get("name").and_then(|v| v.as_str()) {
        parts.push(name.to_string());
    }
    if let Some(desc) = tool.get("description").and_then(|v| v.as_str()) {
        parts.push(desc.to_string());
    }
    if let Some(props) = tool
        .get("inputSchema")
        .and_then(|s| s.get("properties"))
        .and_then(|p| p.as_object())
    {
        for key in props.keys() {
            parts.push(key.clone());
        }
    }
    parts.join(" ").to_lowercase()
}

/// True when any keyword is a substring of the already-lowercased haystack.
fn matches_any(haystack: &str, keywords: &[&str]) -> bool {
    keywords.iter().any(|kw| haystack.contains(kw))
}

/// One classified tool: its name and the tiers it holds.
struct ToolCapability {
    name: String,
    tiers: Vec<CapabilityTier>,
}

/// Run the toxic-flow family over a snapshot.
///
/// `snapshot` is a `tools/list`-style JSON object; only the `tools` array is
/// read, since the capability tiers and the pairing are about what tools can do.
/// Findings come back in rule-ID order: every classified tool's tier tag
/// (SEC-033), then every untrusted-content source flag (SEC-034), then the
/// headline pairing (SEC-035) when the catalog holds both a source and a sink.
///
/// ```
/// use mcptest_core::security::toxic_flow_findings;
/// use serde_json::json;
///
/// let snapshot = json!({ "tools": [
///     { "name": "fetch_url", "description": "Fetch a web page." },
///     { "name": "send_email", "description": "Send an email with credentials." }
/// ] });
/// let findings = toxic_flow_findings(&snapshot);
/// assert!(findings.iter().any(|f| f.rule_id == "SEC-035"));
/// ```
pub fn toxic_flow_findings(snapshot: &Value) -> Vec<SecurityFinding> {
    let classified = classify_snapshot(snapshot);
    let mut findings = Vec::new();
    findings.extend(tier_tags(&classified));
    findings.extend(untrusted_source_flags(&classified));
    findings.extend(pairing(&classified));
    findings
}

/// Classify every tool in the snapshot, keeping only those that matched at
/// least one tier (an unclassified tool carries no capability signal).
fn classify_snapshot(snapshot: &Value) -> Vec<ToolCapability> {
    let Some(arr) = snapshot.get("tools").and_then(|v| v.as_array()) else {
        return Vec::new();
    };
    let mut out = Vec::new();
    for tool in arr {
        let tiers = classify_tool(tool);
        if tiers.is_empty() {
            continue;
        }
        let name = tool
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or("<unnamed>")
            .to_string();
        out.push(ToolCapability { name, tiers });
    }
    out
}

/// Build a toxic-flow finding. These are about a tool name (or the catalog for
/// the pairing), so `definition` carries that name and `kind` is always
/// [`DefinitionKind::Tool`].
fn finding(
    rule_id: &str,
    name: &str,
    severity: SecuritySeverity,
    definition: String,
    message: String,
    evidence: String,
) -> SecurityFinding {
    SecurityFinding {
        rule_id: rule_id.to_string(),
        name: name.to_string(),
        severity,
        category: CheckCategory::Capability,
        definition,
        kind: DefinitionKind::Tool,
        message,
        evidence,
    }
}

// ---------------------------------------------------------------------------
// SEC-033 capability-tier
// ---------------------------------------------------------------------------

/// One informational finding per classified tool, tagging the tiers it holds.
/// This is a posture signal, not a defect: it tells a reviewer what each tool
/// can do so the pairing finding has named context.
fn tier_tags(classified: &[ToolCapability]) -> Vec<SecurityFinding> {
    classified
        .iter()
        .map(|tc| {
            let tiers: Vec<&str> = tc.tiers.iter().map(|t| t.as_str()).collect();
            finding(
                "SEC-033",
                "capability-tier",
                SecuritySeverity::Info,
                tc.name.clone(),
                "tool classified into one or more latent capability tiers".into(),
                tiers.join(", "),
            )
        })
        .collect()
}

// ---------------------------------------------------------------------------
// SEC-034 untrusted-content-source
// ---------------------------------------------------------------------------

/// One low-severity finding per tool that pulls untrusted content. On its own a
/// reader is not an attack, but it is the injection entry point, so it is worth
/// surfacing even when no sink is present.
fn untrusted_source_flags(classified: &[ToolCapability]) -> Vec<SecurityFinding> {
    classified
        .iter()
        .filter(|tc| tc.tiers.contains(&CapabilityTier::UntrustedContentSource))
        .map(|tc| {
            finding(
                "SEC-034",
                "untrusted-content-source",
                SecuritySeverity::Low,
                tc.name.clone(),
                "tool pulls untrusted external content, the entry point for an injection".into(),
                CapabilityTier::UntrustedContentSource.as_str().to_string(),
            )
        })
        .collect()
}

// ---------------------------------------------------------------------------
// SEC-035 toxic-flow-pairing
// ---------------------------------------------------------------------------

/// The headline finding: the catalog holds both an untrusted-content source and
/// an exfil-or-destructive sink, so a prompt injection delivered through the
/// source can drive the sink. One finding per (source, sink) pair, naming both
/// tools as evidence. A sink that exposes sensitive data or runs a remote
/// destructive action is graded high; a local-only destructive sink is medium.
fn pairing(classified: &[ToolCapability]) -> Vec<SecurityFinding> {
    let sources: Vec<&ToolCapability> = classified
        .iter()
        .filter(|tc| tc.tiers.contains(&CapabilityTier::UntrustedContentSource))
        .collect();
    if sources.is_empty() {
        return Vec::new();
    }
    let mut out = Vec::new();
    for source in &sources {
        for sink in classified.iter() {
            // A tool is never its own sink in a pairing.
            if sink.name == source.name {
                continue;
            }
            let sink_tier = sink.tiers.iter().copied().find(|t| t.is_sink());
            let Some(sink_tier) = sink_tier else {
                continue;
            };
            let severity = match sink_tier {
                CapabilityTier::LocalDestructive => SecuritySeverity::Medium,
                _ => SecuritySeverity::High,
            };
            out.push(finding(
                "SEC-035",
                "toxic-flow-pairing",
                severity,
                source.name.clone(),
                "an untrusted-content source and an exfil-or-destructive sink coexist, so an injected instruction can chain them".into(),
                format!(
                    "source {} (untrusted-content-source) -> sink {} ({})",
                    source.name,
                    sink.name,
                    sink_tier.as_str()
                ),
            ));
        }
    }
    out
}

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

    fn tool(name: &str, description: &str) -> Value {
        json!({ "name": name, "description": description })
    }

    #[test]
    fn classify_tags_untrusted_source_on_fetch_keywords() {
        let tiers = classify_tool(&tool("fetch_page", "Fetch a web URL and return its body."));
        assert!(tiers.contains(&CapabilityTier::UntrustedContentSource));
    }

    #[test]
    fn classify_tags_sensitive_exposure_on_credential_keywords() {
        let tiers = classify_tool(&tool(
            "read_vault",
            "Read a secret credential from the vault.",
        ));
        assert!(tiers.contains(&CapabilityTier::SensitiveDataExposure));
    }

    #[test]
    fn classify_tags_workspace_exposure_on_file_keywords() {
        let tiers = classify_tool(&tool("read_file", "Read a file from the repository."));
        assert!(tiers.contains(&CapabilityTier::WorkspaceDataExposure));
    }

    #[test]
    fn classify_tags_destructive_on_delete_and_pay_keywords() {
        let del = classify_tool(&tool("delete_record", "Delete a record by id."));
        assert!(del.contains(&CapabilityTier::Destructive));
        let pay = classify_tool(&tool("send_payment", "Transfer funds and pay an invoice."));
        assert!(pay.contains(&CapabilityTier::Destructive));
    }

    #[test]
    fn classify_tags_local_destructive_on_unlink_keywords() {
        let tiers = classify_tool(&tool(
            "cleanup",
            "Run unlink on a local path to free space.",
        ));
        assert!(tiers.contains(&CapabilityTier::LocalDestructive));
    }

    #[test]
    fn classify_reads_input_schema_property_names() {
        let with_schema = json!({
            "name": "lookup",
            "description": "Look something up.",
            "inputSchema": { "type": "object", "properties": { "url": { "type": "string" } } }
        });
        let tiers = classify_tool(&with_schema);
        assert!(tiers.contains(&CapabilityTier::UntrustedContentSource));
    }

    #[test]
    fn classify_returns_empty_for_a_benign_tool() {
        let tiers = classify_tool(&tool("add", "Add two integers and return the sum."));
        assert!(tiers.is_empty());
    }

    #[test]
    fn classify_dedups_a_tier_matched_by_two_keywords() {
        // Both "delete" and "drop" map to Destructive; it should appear once.
        let tiers = classify_tool(&tool(
            "delete_and_drop",
            "Delete the row and drop the table.",
        ));
        let destructive = tiers
            .iter()
            .filter(|t| **t == CapabilityTier::Destructive)
            .count();
        assert_eq!(destructive, 1);
    }

    #[test]
    fn reader_only_snapshot_flags_no_pairing() {
        let snapshot = json!({ "tools": [
            tool("fetch_url", "Fetch a web page by URL."),
            tool("read_file", "Read a file from the repo."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        assert!(findings.iter().all(|f| f.rule_id != "SEC-035"));
        // The untrusted reader is still flagged on its own.
        assert!(findings.iter().any(|f| f.rule_id == "SEC-034"));
    }

    #[test]
    fn reader_plus_exfil_sink_flags_pairing_with_both_names() {
        let snapshot = json!({ "tools": [
            tool("fetch_url", "Fetch a web page by URL."),
            tool("send_email", "Send an email containing a credential."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        let pair = findings.iter().find(|f| f.rule_id == "SEC-035");
        let pair = match pair {
            Some(pair) => pair,
            None => unreachable!("a source plus a sensitive sink must pair"),
        };
        assert_eq!(pair.name, "toxic-flow-pairing");
        assert_eq!(pair.severity, SecuritySeverity::High);
        assert!(pair.evidence.contains("fetch_url"));
        assert!(pair.evidence.contains("send_email"));
    }

    #[test]
    fn local_destructive_sink_pairs_at_medium_severity() {
        let snapshot = json!({ "tools": [
            tool("scrape", "Crawl and scrape a third-party site."),
            tool("cleanup", "Run unlink on a local file."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        let pair = findings.iter().find(|f| f.rule_id == "SEC-035");
        let pair = match pair {
            Some(pair) => pair,
            None => unreachable!("a source plus a local-destructive sink must pair"),
        };
        assert_eq!(pair.severity, SecuritySeverity::Medium);
    }

    #[test]
    fn benign_tool_yields_no_finding() {
        let snapshot = json!({ "tools": [
            tool("add", "Add two integers."),
            tool("multiply", "Multiply two integers."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        assert!(findings.is_empty());
    }

    #[test]
    fn tier_tag_is_informational_and_lists_every_tier() {
        let snapshot = json!({ "tools": [
            // A single tool that both fetches and deletes carries two tiers.
            tool("fetch_and_delete", "Fetch a URL then delete the local record."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        let tag = findings.iter().find(|f| f.rule_id == "SEC-033");
        let tag = match tag {
            Some(tag) => tag,
            None => unreachable!("a classified tool must carry a tier tag"),
        };
        assert_eq!(tag.severity, SecuritySeverity::Info);
        assert!(tag.evidence.contains("untrusted-content-source"));
        assert!(tag.evidence.contains("destructive"));
    }

    #[test]
    fn a_self_only_source_and_sink_does_not_pair_with_itself() {
        // One tool is both source and sink; with no other tool there is no pair.
        let snapshot = json!({ "tools": [
            tool("fetch_and_delete", "Fetch a URL then delete a record."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        assert!(findings.iter().all(|f| f.rule_id != "SEC-035"));
    }

    #[test]
    fn findings_come_back_in_rule_id_order() {
        let snapshot = json!({ "tools": [
            tool("fetch_url", "Fetch a web page by URL."),
            tool("send_email", "Send an email with a credential."),
        ] });
        let findings = toxic_flow_findings(&snapshot);
        let ids: Vec<&str> = findings.iter().map(|f| f.rule_id.as_str()).collect();
        let first_034 = ids.iter().position(|id| *id == "SEC-034");
        let first_035 = ids.iter().position(|id| *id == "SEC-035");
        let last_033 = ids.iter().rposition(|id| *id == "SEC-033");
        match (last_033, first_034, first_035) {
            (Some(l33), Some(f34), Some(f35)) => {
                assert!(l33 < f34, "all SEC-033 precede SEC-034");
                assert!(f34 < f35, "SEC-034 precedes SEC-035");
            }
            _ => unreachable!("expected all three rule ids to be present"),
        }
    }
}
