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

/// Create a pentest session using the wizard configuration
#[server]
pub async fn create_pentest_session_wizard(
    config_json: 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 config: serde_json::Value =
        serde_json::from_str(&config_json).map_err(|e| ServerFnError::new(e.to_string()))?;
    let client = reqwest::Client::new();
    let resp = client
        .post(&url)
        .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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    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 api_url = format!(
        "{}/api/v1/pentest/lookup-repo?url={}",
        state.agent_api_url, encoded_url
    );
    let resp = reqwest::get(&api_url)
        .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 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 stop_pentest_session(session_id: String) -> Result<(), ServerFnError> {
    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/stop",
        state.agent_api_url
    );
    let client = reqwest::Client::new();
    client
        .post(&url)
        .send()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;
    Ok(())
}

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

#[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 state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;
    let url = format!(
        "{}/api/v1/pentest/sessions/{session_id}/export",
        state.agent_api_url
    );
    let client = reqwest::Client::new();
    let resp = client
        .post(&url)
        .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)
}