compliance-scanner-agent/compliance-mcp/src/tools/pentest.rs

use mongodb::bson::doc;
use rmcp::{model::*, ErrorData as McpError};
use schemars::JsonSchema;
use serde::Deserialize;

use crate::database::Database;

const MAX_LIMIT: i64 = 200;
const DEFAULT_LIMIT: i64 = 50;

fn cap_limit(limit: Option<i64>) -> i64 {
    limit.unwrap_or(DEFAULT_LIMIT).clamp(1, MAX_LIMIT)
}

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

    #[test]
    fn cap_limit_default() {
        assert_eq!(cap_limit(None), DEFAULT_LIMIT);
    }

    #[test]
    fn cap_limit_clamps_high() {
        assert_eq!(cap_limit(Some(1000)), MAX_LIMIT);
    }

    #[test]
    fn cap_limit_clamps_low() {
        assert_eq!(cap_limit(Some(-100)), 1);
        assert_eq!(cap_limit(Some(0)), 1);
    }

    #[test]
    fn cap_limit_normal() {
        assert_eq!(cap_limit(Some(42)), 42);
    }

    #[test]
    fn list_pentest_sessions_params_deserialize() {
        let json = serde_json::json!({
            "target_id": "tgt",
            "status": "running",
            "strategy": "aggressive",
            "limit": 5
        });
        let params: ListPentestSessionsParams = serde_json::from_value(json).unwrap();
        assert_eq!(params.target_id.as_deref(), Some("tgt"));
        assert_eq!(params.status.as_deref(), Some("running"));
        assert_eq!(params.strategy.as_deref(), Some("aggressive"));
        assert_eq!(params.limit, Some(5));
    }

    #[test]
    fn list_pentest_sessions_params_all_optional() {
        let params: ListPentestSessionsParams =
            serde_json::from_value(serde_json::json!({})).unwrap();
        assert!(params.target_id.is_none());
        assert!(params.status.is_none());
        assert!(params.strategy.is_none());
        assert!(params.limit.is_none());
    }

    #[test]
    fn get_pentest_session_params_deserialize() {
        let params: GetPentestSessionParams =
            serde_json::from_value(serde_json::json!({ "id": "abc123" })).unwrap();
        assert_eq!(params.id, "abc123");
    }

    #[test]
    fn get_attack_chain_params_deserialize() {
        let params: GetAttackChainParams =
            serde_json::from_value(serde_json::json!({ "session_id": "s1", "limit": 20 })).unwrap();
        assert_eq!(params.session_id, "s1");
        assert_eq!(params.limit, Some(20));
    }

    #[test]
    fn get_pentest_messages_params_deserialize() {
        let params: GetPentestMessagesParams =
            serde_json::from_value(serde_json::json!({ "session_id": "s2" })).unwrap();
        assert_eq!(params.session_id, "s2");
        assert!(params.limit.is_none());
    }

    #[test]
    fn pentest_stats_params_deserialize() {
        let params: PentestStatsParams =
            serde_json::from_value(serde_json::json!({ "target_id": "t1" })).unwrap();
        assert_eq!(params.target_id.as_deref(), Some("t1"));

        let params2: PentestStatsParams = serde_json::from_value(serde_json::json!({})).unwrap();
        assert!(params2.target_id.is_none());
    }
}

// ── List Pentest Sessions ──────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
pub struct ListPentestSessionsParams {
    /// Filter by target ID
    pub target_id: Option<String>,
    /// Filter by status: running, paused, completed, failed
    pub status: Option<String>,
    /// Filter by strategy: quick, comprehensive, targeted, aggressive, stealth
    pub strategy: Option<String>,
    /// Maximum number of results (default 50, max 200)
    pub limit: Option<i64>,
}

pub async fn list_pentest_sessions(
    db: &Database,
    params: ListPentestSessionsParams,
) -> Result<CallToolResult, McpError> {
    let mut filter = doc! {};
    if let Some(ref target_id) = params.target_id {
        filter.insert("target_id", target_id);
    }
    if let Some(ref status) = params.status {
        filter.insert("status", status);
    }
    if let Some(ref strategy) = params.strategy {
        filter.insert("strategy", strategy);
    }

    let limit = cap_limit(params.limit);

    let mut cursor = db
        .pentest_sessions()
        .find(filter)
        .sort(doc! { "started_at": -1 })
        .limit(limit)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    let mut results = Vec::new();
    while cursor
        .advance()
        .await
        .map_err(|e| McpError::internal_error(format!("cursor error: {e}"), None))?
    {
        let session = cursor
            .deserialize_current()
            .map_err(|e| McpError::internal_error(format!("deserialize error: {e}"), None))?;
        results.push(session);
    }

    let json = serde_json::to_string_pretty(&results)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

// ── Get Pentest Session ────────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
pub struct GetPentestSessionParams {
    /// Pentest session ID (MongoDB ObjectId hex string)
    pub id: String,
}

pub async fn get_pentest_session(
    db: &Database,
    params: GetPentestSessionParams,
) -> Result<CallToolResult, McpError> {
    let oid = bson::oid::ObjectId::parse_str(&params.id)
        .map_err(|e| McpError::invalid_params(format!("invalid id: {e}"), None))?;

    let session = db
        .pentest_sessions()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?
        .ok_or_else(|| McpError::invalid_params("session not found", None))?;

    let json = serde_json::to_string_pretty(&session)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

// ── Get Attack Chain ───────────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
pub struct GetAttackChainParams {
    /// Pentest session ID to get the attack chain for
    pub session_id: String,
    /// Maximum number of nodes (default 50, max 200)
    pub limit: Option<i64>,
}

pub async fn get_attack_chain(
    db: &Database,
    params: GetAttackChainParams,
) -> Result<CallToolResult, McpError> {
    let limit = cap_limit(params.limit);

    let mut cursor = db
        .attack_chain_nodes()
        .find(doc! { "session_id": &params.session_id })
        .sort(doc! { "started_at": 1 })
        .limit(limit)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    let mut results = Vec::new();
    while cursor
        .advance()
        .await
        .map_err(|e| McpError::internal_error(format!("cursor error: {e}"), None))?
    {
        let node = cursor
            .deserialize_current()
            .map_err(|e| McpError::internal_error(format!("deserialize error: {e}"), None))?;
        results.push(node);
    }

    let json = serde_json::to_string_pretty(&results)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

// ── Get Pentest Messages ───────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
pub struct GetPentestMessagesParams {
    /// Pentest session ID
    pub session_id: String,
    /// Maximum number of messages (default 50, max 200)
    pub limit: Option<i64>,
}

pub async fn get_pentest_messages(
    db: &Database,
    params: GetPentestMessagesParams,
) -> Result<CallToolResult, McpError> {
    let limit = cap_limit(params.limit);

    let mut cursor = db
        .pentest_messages()
        .find(doc! { "session_id": &params.session_id })
        .sort(doc! { "created_at": 1 })
        .limit(limit)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    let mut results = Vec::new();
    while cursor
        .advance()
        .await
        .map_err(|e| McpError::internal_error(format!("cursor error: {e}"), None))?
    {
        let msg = cursor
            .deserialize_current()
            .map_err(|e| McpError::internal_error(format!("deserialize error: {e}"), None))?;
        results.push(msg);
    }

    let json = serde_json::to_string_pretty(&results)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

// ── Pentest Stats ──────────────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
pub struct PentestStatsParams {
    /// Filter stats by target ID
    pub target_id: Option<String>,
}

pub async fn pentest_stats(
    db: &Database,
    params: PentestStatsParams,
) -> Result<CallToolResult, McpError> {
    let mut base_filter = doc! {};
    if let Some(ref target_id) = params.target_id {
        base_filter.insert("target_id", target_id);
    }

    // Count running sessions
    let mut running_filter = base_filter.clone();
    running_filter.insert("status", "running");
    let running = db
        .pentest_sessions()
        .count_documents(running_filter)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    // Count total sessions
    let total_sessions = db
        .pentest_sessions()
        .count_documents(base_filter.clone())
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    // Get findings for these sessions — query DAST findings with session_id set
    let mut findings_filter = doc! { "session_id": { "$ne": null } };
    if let Some(ref target_id) = params.target_id {
        findings_filter.insert("target_id", target_id);
    }
    let total_findings = db
        .dast_findings()
        .count_documents(findings_filter.clone())
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    let mut exploitable_filter = findings_filter.clone();
    exploitable_filter.insert("exploitable", true);
    let exploitable = db
        .dast_findings()
        .count_documents(exploitable_filter)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    // Severity counts
    let mut severity = serde_json::Map::new();
    for sev in ["critical", "high", "medium", "low", "info"] {
        let mut sf = findings_filter.clone();
        sf.insert("severity", sev);
        let count = db
            .dast_findings()
            .count_documents(sf)
            .await
            .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;
        severity.insert(sev.to_string(), serde_json::json!(count));
    }

    let summary = serde_json::json!({
        "running_sessions": running,
        "total_sessions": total_sessions,
        "total_findings": total_findings,
        "exploitable_findings": exploitable,
        "severity_distribution": severity,
    });

    let json = serde_json::to_string_pretty(&summary)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}