feat: attack chain DAG visualization, report export, and UI polish
- Add interactive attack chain DAG using vis-network with hierarchical layout, status-colored nodes, risk-based sizing, and click handlers - Add pentest session export API (GET /sessions/:id/export) supporting both JSON and Markdown report formats - Redesign attack chain tab with graph/list toggle views - Add export buttons (MD/JSON) to session header with Blob download - Show exploitable badge and endpoint on finding cards - Add export_pentest_report server function for dashboard Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Sharang Parnerkar
@@ -3,6 +3,7 @@ use std::sync::Arc;
|
||||
use axum::extract::{Extension, Path, Query};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::sse::{Event, Sse};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::Json;
|
||||
use futures_util::stream;
|
||||
use mongodb::bson::doc;
|
||||
@@ -550,3 +551,219 @@ pub async fn get_session_findings(
|
||||
page: Some(params.page),
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct ExportParams {
|
||||
#[serde(default = "default_export_format")]
|
||||
pub format: String,
|
||||
}
|
||||
|
||||
fn default_export_format() -> String {
|
||||
"json".to_string()
|
||||
}
|
||||
|
||||
/// GET /api/v1/pentest/sessions/:id/export?format=json|markdown — Export a session report
|
||||
#[tracing::instrument(skip_all, fields(session_id = %id))]
|
||||
pub async fn export_session_report(
|
||||
Extension(agent): AgentExt,
|
||||
Path(id): Path<String>,
|
||||
Query(params): Query<ExportParams>,
|
||||
) -> Result<axum::response::Response, (StatusCode, String)> {
|
||||
let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
|
||||
|
||||
// Fetch session
|
||||
let session = agent
|
||||
.db
|
||||
.pentest_sessions()
|
||||
.find_one(doc! { "_id": oid })
|
||||
.await
|
||||
.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error: {e}"),
|
||||
)
|
||||
})?
|
||||
.ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
|
||||
|
||||
// Fetch messages
|
||||
let messages: Vec<PentestMessage> = match agent
|
||||
.db
|
||||
.pentest_messages()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "created_at": 1 })
|
||||
.await
|
||||
{
|
||||
Ok(cursor) => collect_cursor_async(cursor).await,
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
|
||||
// Fetch attack chain nodes
|
||||
let nodes: Vec<AttackChainNode> = match agent
|
||||
.db
|
||||
.attack_chain_nodes()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "started_at": 1 })
|
||||
.await
|
||||
{
|
||||
Ok(cursor) => collect_cursor_async(cursor).await,
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
|
||||
// Fetch DAST findings for this session
|
||||
let findings: Vec<DastFinding> = match agent
|
||||
.db
|
||||
.dast_findings()
|
||||
.find(doc! { "session_id": &id })
|
||||
.sort(doc! { "created_at": -1 })
|
||||
.await
|
||||
{
|
||||
Ok(cursor) => collect_cursor_async(cursor).await,
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
|
||||
// Compute severity counts
|
||||
let critical = findings.iter().filter(|f| f.severity.to_string() == "critical").count();
|
||||
let high = findings.iter().filter(|f| f.severity.to_string() == "high").count();
|
||||
let medium = findings.iter().filter(|f| f.severity.to_string() == "medium").count();
|
||||
let low = findings.iter().filter(|f| f.severity.to_string() == "low").count();
|
||||
let info = findings.iter().filter(|f| f.severity.to_string() == "info").count();
|
||||
|
||||
match params.format.as_str() {
|
||||
"markdown" => {
|
||||
let mut md = String::new();
|
||||
md.push_str("# Penetration Test Report\n\n");
|
||||
|
||||
// Executive summary
|
||||
md.push_str("## Executive Summary\n\n");
|
||||
md.push_str(&format!("| Field | Value |\n"));
|
||||
md.push_str("| --- | --- |\n");
|
||||
md.push_str(&format!("| **Session ID** | {} |\n", id));
|
||||
md.push_str(&format!("| **Status** | {} |\n", session.status));
|
||||
md.push_str(&format!("| **Strategy** | {} |\n", session.strategy));
|
||||
md.push_str(&format!("| **Target ID** | {} |\n", session.target_id));
|
||||
md.push_str(&format!(
|
||||
"| **Started** | {} |\n",
|
||||
session.started_at.to_rfc3339()
|
||||
));
|
||||
if let Some(ref completed) = session.completed_at {
|
||||
md.push_str(&format!(
|
||||
"| **Completed** | {} |\n",
|
||||
completed.to_rfc3339()
|
||||
));
|
||||
}
|
||||
md.push_str(&format!(
|
||||
"| **Tool Invocations** | {} |\n",
|
||||
session.tool_invocations
|
||||
));
|
||||
md.push_str(&format!(
|
||||
"| **Success Rate** | {:.1}% |\n",
|
||||
session.success_rate()
|
||||
));
|
||||
md.push('\n');
|
||||
|
||||
// Findings by severity
|
||||
md.push_str("## Findings Summary\n\n");
|
||||
md.push_str(&format!(
|
||||
"| Severity | Count |\n| --- | --- |\n| Critical | {} |\n| High | {} |\n| Medium | {} |\n| Low | {} |\n| Info | {} |\n| **Total** | **{}** |\n\n",
|
||||
critical, high, medium, low, info, findings.len()
|
||||
));
|
||||
|
||||
// Findings table
|
||||
if !findings.is_empty() {
|
||||
md.push_str("## Findings Detail\n\n");
|
||||
md.push_str("| # | Severity | Title | Endpoint | Exploitable |\n");
|
||||
md.push_str("| --- | --- | --- | --- | --- |\n");
|
||||
for (i, f) in findings.iter().enumerate() {
|
||||
md.push_str(&format!(
|
||||
"| {} | {} | {} | {} {} | {} |\n",
|
||||
i + 1,
|
||||
f.severity,
|
||||
f.title,
|
||||
f.method,
|
||||
f.endpoint,
|
||||
if f.exploitable { "Yes" } else { "No" },
|
||||
));
|
||||
}
|
||||
md.push('\n');
|
||||
}
|
||||
|
||||
// Attack chain timeline
|
||||
if !nodes.is_empty() {
|
||||
md.push_str("## Attack Chain Timeline\n\n");
|
||||
md.push_str("| # | Tool | Status | Findings | Reasoning |\n");
|
||||
md.push_str("| --- | --- | --- | --- | --- |\n");
|
||||
for (i, node) in nodes.iter().enumerate() {
|
||||
let reasoning_short = if node.llm_reasoning.len() > 80 {
|
||||
format!("{}...", &node.llm_reasoning[..80])
|
||||
} else {
|
||||
node.llm_reasoning.clone()
|
||||
};
|
||||
md.push_str(&format!(
|
||||
"| {} | {} | {} | {} | {} |\n",
|
||||
i + 1,
|
||||
node.tool_name,
|
||||
format!("{:?}", node.status).to_lowercase(),
|
||||
node.findings_produced.len(),
|
||||
reasoning_short,
|
||||
));
|
||||
}
|
||||
md.push('\n');
|
||||
}
|
||||
|
||||
// Statistics
|
||||
md.push_str("## Statistics\n\n");
|
||||
md.push_str(&format!("- **Total Findings:** {}\n", findings.len()));
|
||||
md.push_str(&format!("- **Exploitable Findings:** {}\n", session.exploitable_count));
|
||||
md.push_str(&format!("- **Attack Chain Steps:** {}\n", nodes.len()));
|
||||
md.push_str(&format!("- **Messages Exchanged:** {}\n", messages.len()));
|
||||
md.push_str(&format!("- **Tool Invocations:** {}\n", session.tool_invocations));
|
||||
md.push_str(&format!("- **Tool Success Rate:** {:.1}%\n", session.success_rate()));
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
[
|
||||
(axum::http::header::CONTENT_TYPE, "text/markdown; charset=utf-8"),
|
||||
],
|
||||
md,
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
_ => {
|
||||
// JSON format
|
||||
let report = serde_json::json!({
|
||||
"session": {
|
||||
"id": id,
|
||||
"target_id": session.target_id,
|
||||
"repo_id": session.repo_id,
|
||||
"status": session.status,
|
||||
"strategy": session.strategy,
|
||||
"started_at": session.started_at.to_rfc3339(),
|
||||
"completed_at": session.completed_at.map(|d| d.to_rfc3339()),
|
||||
"tool_invocations": session.tool_invocations,
|
||||
"tool_successes": session.tool_successes,
|
||||
"success_rate": session.success_rate(),
|
||||
"findings_count": session.findings_count,
|
||||
"exploitable_count": session.exploitable_count,
|
||||
},
|
||||
"findings": findings,
|
||||
"attack_chain": nodes,
|
||||
"messages": messages,
|
||||
"summary": {
|
||||
"total_findings": findings.len(),
|
||||
"severity_distribution": {
|
||||
"critical": critical,
|
||||
"high": high,
|
||||
"medium": medium,
|
||||
"low": low,
|
||||
"info": info,
|
||||
},
|
||||
"attack_chain_steps": nodes.len(),
|
||||
"messages_exchanged": messages.len(),
|
||||
},
|
||||
});
|
||||
|
||||
Ok(Json(report).into_response())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -128,6 +128,10 @@ pub fn build_router() -> Router {
|
||||
"/api/v1/pentest/sessions/{id}/findings",
|
||||
get(handlers::pentest::get_session_findings),
|
||||
)
|
||||
.route(
|
||||
"/api/v1/pentest/sessions/{id}/export",
|
||||
get(handlers::pentest::export_session_report),
|
||||
)
|
||||
.route("/api/v1/pentest/stats", get(handlers::pentest::pentest_stats))
|
||||
// Webhook endpoints (proxied through dashboard)
|
||||
.route(
|
||||
|
||||
Reference in New Issue
Block a user