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 c0f9ba467c
commit 85ceef7e1f
6 changed files with 393 additions and 21 deletions

View File

@@ -102,18 +102,9 @@ pub async fn create_session(
let target_clone = target.clone();
tokio::spawn(async move {
let orchestrator = PentestOrchestrator::new(llm, db);
if let Err(e) = orchestrator
.run_session(&session_clone, &target_clone, &initial_message)
.await
{
tracing::error!(
"Pentest orchestrator failed for session {}: {e}",
session_clone
.id
.map(|oid| oid.to_hex())
.unwrap_or_default()
);
}
orchestrator
.run_session_guarded(&session_clone, &target_clone, &initial_message)
.await;
});
Ok(Json(ApiResponse {
@@ -254,9 +245,9 @@ pub async fn send_message(
let message = req.message.clone();
tokio::spawn(async move {
let orchestrator = PentestOrchestrator::new(llm, db);
if let Err(e) = orchestrator.run_session(&session, &target, &message).await {
tracing::error!("Pentest orchestrator failed for session {session_id}: {e}");
}
orchestrator
.run_session_guarded(&session, &target, &message)
.await;
});
Ok(Json(ApiResponse {
@@ -474,7 +465,6 @@ pub async fn pentest_stats(
};
// Severity distribution from pentest-related DAST findings
let pentest_filter = doc! { "session_id": { "$exists": true, "$ne": null } };
let critical = db
.dast_findings()
.count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "critical" })
@@ -501,8 +491,6 @@ pub async fn pentest_stats(
.await
.unwrap_or(0) as u32;
let _ = pentest_filter; // used above inline
Ok(Json(ApiResponse {
data: PentestStats {
running_sessions,