//! Capability-negotiation integrity check for the 2026-07-28 stateless model
//! (PROTOAMP, AttestMCP; arXiv:2601.17549).
//!
//! Why this is a separate module: the surface and namespace lanes inspect a
//! static catalog. This lane inspects the *handshake* under the 2026-07-28
//! stateless transport, where there is no long-lived session, so every request
//! re-declares the capabilities it claims. "Breaking the Protocol" shows the
//! absence of capability attestation lets a server silently widen what it claims
//! across requests, escalating from the capability set the client negotiated to
//! a broader one the client never agreed to. AttestMCP cut attack success from
//! 52.8 to 12.4 percent by checking exactly this: that per-request capability
//! claims stay consistent with the negotiated baseline and are never widened.
//!
//! The check is relational across requests, so like the other relational lanes
//! it lives outside the per-[`Definition`] engine as a standalone function. It
//! takes the negotiated baseline (the capability set the client agreed at
//! `initialize`) and the per-request claims captured during the run, and flags a
//! request that claims a capability outside the baseline (a silent widening).
//!
//! Detection is an objective set comparison, not an LLM judge (ADR 0040): a claim
//! is a widening exactly when it is present in a request but absent from the
//! negotiated baseline. No model decides the verdict.
//!
//! [`Definition`]: crate::security::Definition

use std::collections::BTreeSet;

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

/// The protocol target this check is scoped to. Capability negotiation is a
/// stateless-model concern, so the check only applies to the 2026-07-28 target.
pub const STATELESS_TARGET: &str = "2026-07-28";

/// The capability claims a single request carried.
///
/// Under the stateless model each request restates its capabilities, so the
/// check captures one of these per request. `request_id` labels the request in
/// findings; `capabilities` is the set of capability keys the request claimed
/// (for example `tools`, `resources/subscribe`, `elicitation`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequestCapabilityClaim {
    /// A stable label for the request, used in finding evidence.
    pub request_id: String,
    /// The capability keys this request claimed.
    pub capabilities: Vec<String>,
}

/// The negotiated baseline plus the per-request claims captured during a run.
///
/// `target` is the protocol version the suite ran against; the check is a no-op
/// unless it is the stateless [`STATELESS_TARGET`]. `negotiated` is the
/// capability set agreed at `initialize`. `requests` is every later request's
/// restated claims.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct NegotiationTrace {
    /// The protocol target the run used, for example `2026-07-28`.
    pub target: String,
    /// The capability set negotiated at `initialize`, the agreed baseline.
    pub negotiated: Vec<String>,
    /// Each subsequent request's restated capability claims.
    pub requests: Vec<RequestCapabilityClaim>,
}

/// Run the capability-negotiation integrity check over a run's negotiation trace.
///
/// The check applies only under the 2026-07-28 stateless target; for any other
/// target it returns no findings, since pre-stateless sessions negotiate
/// capabilities once and do not restate them per request. Under the stateless
/// target, every request whose claimed capability set is not a subset of the
/// negotiated baseline fires SEC-061: the request silently widened its claims
/// beyond what the client agreed. One finding per widening request, naming the
/// extra capabilities.
///
/// ```
/// use mcptest_core::security::{capability_negotiation_findings, NegotiationTrace, RequestCapabilityClaim};
///
/// let trace = NegotiationTrace {
///     target: "2026-07-28".into(),
///     negotiated: vec!["tools".into()],
///     requests: vec![RequestCapabilityClaim {
///         request_id: "req-2".into(),
///         capabilities: vec!["tools".into(), "elicitation".into()],
///     }],
/// };
/// let findings = capability_negotiation_findings(&trace);
/// assert!(findings.iter().any(|f| f.rule_id == "SEC-061"));
/// ```
pub fn capability_negotiation_findings(trace: &NegotiationTrace) -> Vec<SecurityFinding> {
    // Pre-stateless targets negotiate once; the per-request integrity rule does
    // not apply, so there is nothing to flag.
    if trace.target != STATELESS_TARGET {
        return Vec::new();
    }
    let baseline: BTreeSet<&str> = trace.negotiated.iter().map(String::as_str).collect();
    let mut out = Vec::new();
    for request in &trace.requests {
        let widened: Vec<&str> = request
            .capabilities
            .iter()
            .map(String::as_str)
            .filter(|cap| !baseline.contains(cap))
            .collect();
        if !widened.is_empty() {
            out.push(finding(&request.request_id, &widened));
        }
    }
    out
}

/// Build the SEC-061 finding for one request that widened its capability claims.
fn finding(request_id: &str, widened: &[&str]) -> SecurityFinding {
    SecurityFinding {
        rule_id: "SEC-061".to_string(),
        name: "capability-negotiation-widening".to_string(),
        severity: SecuritySeverity::High,
        category: CheckCategory::Auth,
        definition: request_id.to_string(),
        kind: DefinitionKind::Tool,
        message:
            "a request claimed a capability outside the negotiated baseline, a silent widening of the agreed capability set"
                .to_string(),
        evidence: format!("request {request_id} widened: {}", widened.join(", ")),
    }
}

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

    fn claim(id: &str, caps: &[&str]) -> RequestCapabilityClaim {
        RequestCapabilityClaim {
            request_id: id.to_string(),
            capabilities: caps.iter().map(|c| c.to_string()).collect(),
        }
    }

    #[test]
    fn widened_capability_under_stateless_target_fires_at_high() {
        let trace = NegotiationTrace {
            target: STATELESS_TARGET.to_string(),
            negotiated: vec!["tools".into(), "resources".into()],
            requests: vec![claim("req-2", &["tools", "elicitation"])],
        };
        let found = capability_negotiation_findings(&trace);
        assert_eq!(found.len(), 1);
        assert_eq!(found[0].rule_id, "SEC-061");
        assert_eq!(found[0].name, "capability-negotiation-widening");
        assert_eq!(found[0].severity, SecuritySeverity::High);
        assert_eq!(found[0].category, CheckCategory::Auth);
        assert_eq!(found[0].definition, "req-2");
        assert!(found[0].evidence.contains("elicitation"));
    }

    #[test]
    fn consistent_claims_within_the_baseline_pass() {
        let trace = NegotiationTrace {
            target: STATELESS_TARGET.to_string(),
            negotiated: vec!["tools".into(), "resources".into(), "elicitation".into()],
            requests: vec![
                claim("req-2", &["tools"]),
                claim("req-3", &["tools", "resources"]),
                claim("req-4", &["resources", "elicitation"]),
            ],
        };
        let found = capability_negotiation_findings(&trace);
        assert!(found.is_empty());
    }

    #[test]
    fn narrowing_below_the_baseline_is_not_a_finding() {
        // Claiming fewer capabilities than negotiated is safe, never widening.
        let trace = NegotiationTrace {
            target: STATELESS_TARGET.to_string(),
            negotiated: vec!["tools".into(), "resources".into()],
            requests: vec![claim("req-2", &[])],
        };
        let found = capability_negotiation_findings(&trace);
        assert!(found.is_empty());
    }

    #[test]
    fn each_widening_request_gets_its_own_finding() {
        let trace = NegotiationTrace {
            target: STATELESS_TARGET.to_string(),
            negotiated: vec!["tools".into()],
            requests: vec![
                claim("req-2", &["tools", "sampling"]),
                claim("req-3", &["tools"]),
                claim("req-4", &["tools", "roots"]),
            ],
        };
        let found = capability_negotiation_findings(&trace);
        assert_eq!(found.len(), 2);
        assert_eq!(found[0].definition, "req-2");
        assert_eq!(found[1].definition, "req-4");
    }

    #[test]
    fn non_stateless_target_is_a_no_op_even_when_widened() {
        // The same widening under an older session-based target does not apply.
        let trace = NegotiationTrace {
            target: "2025-06-18".to_string(),
            negotiated: vec!["tools".into()],
            requests: vec![claim("req-2", &["tools", "elicitation"])],
        };
        let found = capability_negotiation_findings(&trace);
        assert!(found.is_empty());
    }

    #[test]
    fn report_lists_every_extra_capability_when_several_are_widened() {
        let trace = NegotiationTrace {
            target: STATELESS_TARGET.to_string(),
            negotiated: vec!["tools".into()],
            requests: vec![claim("req-2", &["tools", "sampling", "roots"])],
        };
        let found = capability_negotiation_findings(&trace);
        assert_eq!(found.len(), 1);
        assert!(found[0].evidence.contains("sampling"));
        assert!(found[0].evidence.contains("roots"));
    }
}
