refactor: modularize codebase and add 404 unit tests (#13)

2026-03-13 08:03:45 +00:00
parent acc5b86aa4
commit 3bb690e5bb
89 changed files with 11884 additions and 6046 deletions
@@ -0,0 +1,102 @@
+use std::sync::Arc;
+
+use axum::extract::Extension;
+use axum::http::StatusCode;
+use axum::Json;
+use mongodb::bson::doc;
+
+use compliance_core::models::pentest::*;
+
+use crate::agent::ComplianceAgent;
+
+use super::super::dto::{collect_cursor_async, ApiResponse};
+
+type AgentExt = Extension<Arc<ComplianceAgent>>;
+
+/// GET /api/v1/pentest/stats — Aggregated pentest statistics
+#[tracing::instrument(skip_all)]
+pub async fn pentest_stats(
+    Extension(agent): AgentExt,
+) -> Result<Json<ApiResponse<PentestStats>>, StatusCode> {
+    let db = &agent.db;
+
+    let running_sessions = db
+        .pentest_sessions()
+        .count_documents(doc! { "status": "running" })
+        .await
+        .unwrap_or(0) as u32;
+
+    // Count DAST findings from pentest sessions
+    let total_vulnerabilities = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null } })
+        .await
+        .unwrap_or(0) as u32;
+
+    // Aggregate tool invocations from all sessions
+    let sessions: Vec<PentestSession> = match db.pentest_sessions().find(doc! {}).await {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(_) => Vec::new(),
+    };
+
+    let total_tool_invocations: u32 = sessions.iter().map(|s| s.tool_invocations).sum();
+    let total_successes: u32 = sessions.iter().map(|s| s.tool_successes).sum();
+    let tool_success_rate = if total_tool_invocations == 0 {
+        100.0
+    } else {
+        (total_successes as f64 / total_tool_invocations as f64) * 100.0
+    };
+
+    // Severity distribution from pentest-related DAST findings
+    let critical = db
+        .dast_findings()
+        .count_documents(
+            doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "critical" },
+        )
+        .await
+        .unwrap_or(0) as u32;
+    let high = db
+        .dast_findings()
+        .count_documents(
+            doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "high" },
+        )
+        .await
+        .unwrap_or(0) as u32;
+    let medium = db
+        .dast_findings()
+        .count_documents(
+            doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "medium" },
+        )
+        .await
+        .unwrap_or(0) as u32;
+    let low = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "low" })
+        .await
+        .unwrap_or(0) as u32;
+    let info = db
+        .dast_findings()
+        .count_documents(
+            doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "info" },
+        )
+        .await
+        .unwrap_or(0) as u32;
+
+    Ok(Json(ApiResponse {
+        data: PentestStats {
+            running_sessions,
+            total_vulnerabilities,
+            total_tool_invocations,
+            tool_success_rate,
+            severity_distribution: SeverityDistribution {
+                critical,
+                high,
+                medium,
+                low,
+                info,
+            },
+        },
+        total: None,
+        page: None,
+    }))
+}