feat: add pentest MCP tools, session timeout, and error recovery

Add 5 MCP tools for querying pentest sessions, attack chains, messages,
and stats. Add session timeout (30min) and automatic failure marking
with run_session_guarded wrapper.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Sharang Parnerkar
2026-03-11 19:49:39 +01:00
parent 03d8e16e13
commit ad9036e5ad
6 changed files with 393 additions and 21 deletions

View File

@@ -3,7 +3,7 @@ use rmcp::{
};
use crate::database::Database;
use crate::tools::{dast, findings, sbom};
use crate::tools::{dast, findings, pentest, sbom};
pub struct ComplianceMcpServer {
db: Database,
@@ -89,6 +89,54 @@ impl ComplianceMcpServer {
) -> Result<CallToolResult, rmcp::ErrorData> {
dast::dast_scan_summary(&self.db, params).await
}
// ── Pentest ─────────────────────────────────────────────
#[tool(
description = "List AI pentest sessions with optional filters for target, status, and strategy"
)]
async fn list_pentest_sessions(
&self,
Parameters(params): Parameters<pentest::ListPentestSessionsParams>,
) -> Result<CallToolResult, rmcp::ErrorData> {
pentest::list_pentest_sessions(&self.db, params).await
}
#[tool(description = "Get a single AI pentest session by its ID")]
async fn get_pentest_session(
&self,
Parameters(params): Parameters<pentest::GetPentestSessionParams>,
) -> Result<CallToolResult, rmcp::ErrorData> {
pentest::get_pentest_session(&self.db, params).await
}
#[tool(
description = "Get the attack chain DAG for a pentest session showing each tool invocation, its reasoning, and results"
)]
async fn get_attack_chain(
&self,
Parameters(params): Parameters<pentest::GetAttackChainParams>,
) -> Result<CallToolResult, rmcp::ErrorData> {
pentest::get_attack_chain(&self.db, params).await
}
#[tool(description = "Get chat messages from a pentest session")]
async fn get_pentest_messages(
&self,
Parameters(params): Parameters<pentest::GetPentestMessagesParams>,
) -> Result<CallToolResult, rmcp::ErrorData> {
pentest::get_pentest_messages(&self.db, params).await
}
#[tool(
description = "Get aggregated pentest statistics including running sessions, vulnerability counts, and severity distribution"
)]
async fn pentest_stats(
&self,
Parameters(params): Parameters<pentest::PentestStatsParams>,
) -> Result<CallToolResult, rmcp::ErrorData> {
pentest::pentest_stats(&self.db, params).await
}
}
#[tool_handler]
@@ -101,7 +149,7 @@ impl ServerHandler for ComplianceMcpServer {
.build(),
server_info: Implementation::from_build_env(),
instructions: Some(
"Compliance Scanner MCP server. Query security findings, SBOM data, and DAST results."
"Compliance Scanner MCP server. Query security findings, SBOM data, DAST results, and AI pentest sessions."
.to_string(),
),
}