compliance-scanner-agent/compliance-agent/src/api/handlers/pentest_handlers/stats.rs

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,
    }))
}