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

use std::sync::Arc;

use axum::extract::{Extension, Path};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::Json;
use mongodb::bson::doc;
use serde::Deserialize;

use compliance_core::models::dast::DastFinding;
use compliance_core::models::pentest::*;

use crate::agent::ComplianceAgent;

use super::super::dto::collect_cursor_async;

type AgentExt = Extension<Arc<ComplianceAgent>>;

#[derive(Deserialize)]
pub struct ExportBody {
    pub password: String,
    /// Requester display name (from auth)
    #[serde(default)]
    pub requester_name: String,
    /// Requester email (from auth)
    #[serde(default)]
    pub requester_email: String,
}

/// POST /api/v1/pentest/sessions/:id/export — Export an encrypted pentest report archive
#[tracing::instrument(skip_all, fields(session_id = %id))]
pub async fn export_session_report(
    Extension(agent): AgentExt,
    Path(id): Path<String>,
    Json(body): Json<ExportBody>,
) -> 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()))?;

    if body.password.len() < 8 {
        return Err((
            StatusCode::BAD_REQUEST,
            "Password must be at least 8 characters".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()))?;

    // Resolve target name
    let target = if let Ok(tid) = mongodb::bson::oid::ObjectId::parse_str(&session.target_id) {
        agent
            .db
            .dast_targets()
            .find_one(doc! { "_id": tid })
            .await
            .ok()
            .flatten()
    } else {
        None
    };
    let target_name = target
        .as_ref()
        .map(|t| t.name.clone())
        .unwrap_or_else(|| "Unknown Target".to_string());
    let target_url = target
        .as_ref()
        .map(|t| t.base_url.clone())
        .unwrap_or_default();

    // 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! { "severity": -1, "created_at": -1 })
        .await
    {
        Ok(cursor) => collect_cursor_async(cursor).await,
        Err(_) => Vec::new(),
    };

    let ctx = crate::pentest::report::ReportContext {
        session,
        target_name,
        target_url,
        findings,
        attack_chain: nodes,
        requester_name: if body.requester_name.is_empty() {
            "Unknown".to_string()
        } else {
            body.requester_name
        },
        requester_email: body.requester_email,
    };

    let report = crate::pentest::generate_encrypted_report(&ctx, &body.password)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;

    let response = serde_json::json!({
        "archive_base64": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &report.archive),
        "sha256": report.sha256,
        "filename": format!("pentest-report-{id}.zip"),
    });

    Ok(Json(response).into_response())
}