feat: AI-driven automated penetration testing (#12)

compliance-mcp/src/tools/mod.rs

@@ -1,3 +1,4 @@
 pub mod dast;
 pub mod findings;
+pub mod pentest;
 pub mod sbom;

compliance-mcp/src/tools/pentest.rs

@@ -0,0 +1,261 @@
+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)
+}
+
+// ── 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)]))
+}