compliance-scanner-agent/compliance-dashboard/src/infrastructure/pentest.rs

use dioxus::prelude::*;
use serde::{Deserialize, Serialize};

use super::dast::DastFindingsResponse;

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PentestSessionsResponse {
    pub data: Vec<serde_json::Value>,
    pub total: Option<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PentestSessionResponse {
    pub data: serde_json::Value,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PentestMessagesResponse {
    pub data: Vec<serde_json::Value>,
    pub total: Option<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PentestStatsResponse {
    pub data: serde_json::Value,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct AttackChainResponse {
    pub data: Vec<serde_json::Value>,
}

#[server]
pub async fn fetch_pentest_sessions() -> Result<PentestSessionsResponse, ServerFnError> {
    // Fetch sessions
    let resp = super::agent_client::agent_get("/api/v1/pentest/sessions")
        .await?
        .send()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    let mut body: PentestSessionsResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    // Fetch DAST targets to resolve target names
    if let Ok(tresp_builder) = super::agent_client::agent_get("/api/v1/dast/targets").await {
        if let Ok(tresp) = tresp_builder.send().await {
            if let Ok(tbody) = tresp.json::<serde_json::Value>().await {
                let targets = tbody.get("data").and_then(|v| v.as_array());
                if let Some(targets) = targets {
                    // Build target_id -> name lookup
                    let target_map: std::collections::HashMap<String, String> = targets
                        .iter()
                        .filter_map(|t| {
                            let id = t.get("_id")?.get("$oid")?.as_str()?.to_string();
                            let name = t.get("name")?.as_str()?.to_string();
                            Some((id, name))
                        })
                        .collect();

                    // Enrich sessions with target_name
                    for session in body.data.iter_mut() {
                        if let Some(tid) = session.get("target_id").and_then(|v| v.as_str()) {
                            if let Some(name) = target_map.get(tid) {
                                session.as_object_mut().map(|obj| {
                                    obj.insert(
                                        "target_name".to_string(),
                                        serde_json::Value::String(name.clone()),
                                    )
                                });
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(body)
}

#[server]
pub async fn fetch_pentest_session(id: String) -> Result<PentestSessionResponse, ServerFnError> {
    let resp = super::agent_client::agent_get(&format!("/api/v1/pentest/sessions/{id}"))
        .await?
        .send()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    let mut body: PentestSessionResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    // Resolve target name from targets list
    if let Some(tid) = body.data.get("target_id").and_then(|v| v.as_str()) {
        if let Ok(tresp_builder) = super::agent_client::agent_get("/api/v1/dast/targets").await {
            if let Ok(tresp) = tresp_builder.send().await {
                if let Ok(tbody) = tresp.json::<serde_json::Value>().await {
                    if let Some(targets) = tbody.get("data").and_then(|v| v.as_array()) {
                        for t in targets {
                            let t_id = t
                                .get("_id")
                                .and_then(|v| v.get("$oid"))
                                .and_then(|v| v.as_str())
                                .unwrap_or("");
                            if t_id == tid {
                                if let Some(name) = t.get("name").and_then(|v| v.as_str()) {
                                    body.data.as_object_mut().map(|obj| {
                                        obj.insert(
                                            "target_name".to_string(),
                                            serde_json::Value::String(name.to_string()),
                                        )
                                    });
                                }
                                break;
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(body)
}

#[server]
pub async fn fetch_pentest_messages(
    session_id: String,
) -> Result<PentestMessagesResponse, ServerFnError> {
    let resp =
        super::agent_client::agent_get(&format!("/api/v1/pentest/sessions/{session_id}/messages"))
            .await?
            .send()
            .await
            .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: PentestMessagesResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

#[server]
pub async fn fetch_pentest_stats() -> Result<PentestStatsResponse, ServerFnError> {
    let resp = super::agent_client::agent_get("/api/v1/pentest/stats")
        .await?
        .send()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: PentestStatsResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

#[server]
pub async fn fetch_attack_chain(session_id: String) -> Result<AttackChainResponse, ServerFnError> {
    let resp = super::agent_client::agent_get(&format!(
        "/api/v1/pentest/sessions/{session_id}/attack-chain"
    ))
    .await?
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: AttackChainResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

#[server]
pub async fn create_pentest_session(
    target_id: String,
    strategy: String,
    message: String,
) -> Result<PentestSessionResponse, ServerFnError> {
    let resp =
        super::agent_client::agent_request(reqwest::Method::POST, "/api/v1/pentest/sessions")
            .await?
            .json(&serde_json::json!({
                "target_id": target_id,
                "strategy": strategy,
                "message": message,
            }))
            .send()
            .await
            .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: PentestSessionResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

/// Create a pentest session using the wizard configuration
#[server]
pub async fn create_pentest_session_wizard(
    config_json: String,
) -> Result<PentestSessionResponse, ServerFnError> {
    let config: serde_json::Value =
        serde_json::from_str(&config_json).map_err(|e| ServerFnError::new(e.to_string()))?;
    let resp =
        super::agent_client::agent_request(reqwest::Method::POST, "/api/v1/pentest/sessions")
            .await?
            .json(&serde_json::json!({ "config": config }))
            .send()
            .await
            .map_err(|e| ServerFnError::new(e.to_string()))?;
    if !resp.status().is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!(
            "Failed to create session: {text}"
        )));
    }
    let body: PentestSessionResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

/// Look up a tracked repository by its git URL
#[server]
pub async fn lookup_repo_by_url(url: String) -> Result<serde_json::Value, ServerFnError> {
    let encoded_url: String = url
        .bytes()
        .flat_map(|b| {
            if b.is_ascii_alphanumeric() || b == b'-' || b == b'_' || b == b'.' || b == b'~' {
                vec![b as char]
            } else {
                format!("%{:02X}", b).chars().collect()
            }
        })
        .collect();
    let resp =
        super::agent_client::agent_get(&format!("/api/v1/pentest/lookup-repo?url={encoded_url}"))
            .await?
            .send()
            .await
            .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body.get("data").cloned().unwrap_or(serde_json::Value::Null))
}

#[server]
pub async fn send_pentest_message(
    session_id: String,
    message: String,
) -> Result<PentestMessagesResponse, ServerFnError> {
    let resp = super::agent_client::agent_request(
        reqwest::Method::POST,
        &format!("/api/v1/pentest/sessions/{session_id}/chat"),
    )
    .await?
    .json(&serde_json::json!({
        "message": message,
    }))
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: PentestMessagesResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

#[server]
pub async fn stop_pentest_session(session_id: String) -> Result<(), ServerFnError> {
    super::agent_client::agent_request(
        reqwest::Method::POST,
        &format!("/api/v1/pentest/sessions/{session_id}/stop"),
    )
    .await?
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(())
}

#[server]
pub async fn pause_pentest_session(session_id: String) -> Result<(), ServerFnError> {
    let resp = super::agent_client::agent_request(
        reqwest::Method::POST,
        &format!("/api/v1/pentest/sessions/{session_id}/pause"),
    )
    .await?
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    if !resp.status().is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!("Pause failed: {text}")));
    }
    Ok(())
}

#[server]
pub async fn resume_pentest_session(session_id: String) -> Result<(), ServerFnError> {
    let resp = super::agent_client::agent_request(
        reqwest::Method::POST,
        &format!("/api/v1/pentest/sessions/{session_id}/resume"),
    )
    .await?
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    if !resp.status().is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!("Resume failed: {text}")));
    }
    Ok(())
}

#[server]
pub async fn fetch_pentest_findings(
    session_id: String,
) -> Result<DastFindingsResponse, ServerFnError> {
    let resp =
        super::agent_client::agent_get(&format!("/api/v1/pentest/sessions/{session_id}/findings"))
            .await?
            .send()
            .await
            .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body: DastFindingsResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ExportReportResponse {
    pub archive_base64: String,
    pub sha256: String,
    pub filename: String,
}

#[server]
pub async fn export_pentest_report(
    session_id: String,
    password: String,
    requester_name: String,
    requester_email: String,
) -> Result<ExportReportResponse, ServerFnError> {
    let resp = super::agent_client::agent_request(
        reqwest::Method::POST,
        &format!("/api/v1/pentest/sessions/{session_id}/export"),
    )
    .await?
    .json(&serde_json::json!({
        "password": password,
        "requester_name": requester_name,
        "requester_email": requester_email,
    }))
    .send()
    .await
    .map_err(|e| ServerFnError::new(e.to_string()))?;
    if !resp.status().is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!("Export failed: {text}")));
    }
    let body: ExportReportResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}