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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user