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> {
    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    // Fetch sessions
    let url = format!("{}/api/v1/pentest/sessions", state.agent_api_url);
    let resp = reqwest::get(&url)
        .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
    let targets_url = format!("{}/api/v1/dast/targets", state.agent_api_url);
    if let Ok(tresp) = reqwest::get(&targets_url).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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!("{}/api/v1/pentest/sessions/{id}", state.agent_api_url);
    let resp = reqwest::get(&url)
        .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()) {
        let targets_url = format!("{}/api/v1/dast/targets", state.agent_api_url);
        if let Ok(tresp) = reqwest::get(&targets_url).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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/messages",
        state.agent_api_url
    );
    let resp = reqwest::get(&url)
        .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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!("{}/api/v1/pentest/stats", state.agent_api_url);
    let resp = reqwest::get(&url)
        .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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/attack-chain",
        state.agent_api_url
    );
    let resp = reqwest::get(&url)
        .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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!("{}/api/v1/pentest/sessions", state.agent_api_url);
    let client = reqwest::Client::new();
    let resp = client
        .post(&url)
        .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)
}

#[server]
pub async fn send_pentest_message(
    session_id: String,
    message: String,
) -> Result<PentestMessagesResponse, ServerFnError> {
    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/chat",
        state.agent_api_url
    );
    let client = reqwest::Client::new();
    let resp = client
        .post(&url)
        .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 fetch_pentest_findings(
    session_id: String,
) -> Result<DastFindingsResponse, ServerFnError> {
    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/findings",
        state.agent_api_url
    );
    let resp = reqwest::get(&url)
        .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)
}

#[server]
pub async fn export_pentest_report(
    session_id: String,
    format: String,
) -> Result<String, ServerFnError> {
    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/export?format={format}",
        state.agent_api_url
    );
    let resp = reqwest::get(&url)
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    let body = resp
        .text()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(body)
}