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

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

use compliance_core::models::McpServerConfig;

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

#[server]
pub async fn fetch_mcp_servers() -> Result<McpServersResponse, ServerFnError> {
    use mongodb::bson::doc;

    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let mut cursor = state
        .db
        .mcp_servers()
        .find(doc! {})
        .sort(doc! { "created_at": -1 })
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    let mut data = Vec::new();
    while cursor
        .advance()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?
    {
        let server = cursor
            .deserialize_current()
            .map_err(|e| ServerFnError::new(e.to_string()))?;
        data.push(server);
    }

    Ok(McpServersResponse { data })
}

#[server]
pub async fn add_mcp_server(
    name: String,
    endpoint_url: String,
    transport: String,
    port: String,
    description: String,
    mongodb_uri: String,
    mongodb_database: String,
) -> Result<(), ServerFnError> {
    use chrono::Utc;
    use compliance_core::models::{McpServerStatus, McpTransport};

    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let transport = match transport.as_str() {
        "http" => McpTransport::Http,
        _ => McpTransport::Stdio,
    };

    let port_num: Option<u16> = port.parse().ok();

    // Generate a random access token
    let token = format!("mcp_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));

    let all_tools = vec![
        "list_findings".to_string(),
        "get_finding".to_string(),
        "findings_summary".to_string(),
        "list_sbom_packages".to_string(),
        "sbom_vuln_report".to_string(),
        "list_dast_findings".to_string(),
        "dast_scan_summary".to_string(),
    ];

    let now = Utc::now();
    let server = McpServerConfig {
        id: None,
        name,
        endpoint_url,
        transport,
        port: port_num,
        status: McpServerStatus::Stopped,
        access_token: token,
        tools_enabled: all_tools,
        description: if description.is_empty() {
            None
        } else {
            Some(description)
        },
        mongodb_uri: if mongodb_uri.is_empty() {
            None
        } else {
            Some(mongodb_uri)
        },
        mongodb_database: if mongodb_database.is_empty() {
            None
        } else {
            Some(mongodb_database)
        },
        created_at: now,
        updated_at: now,
    };

    state
        .db
        .mcp_servers()
        .insert_one(server)
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    Ok(())
}

/// Probe each MCP server's health endpoint and update status in MongoDB.
#[server]
pub async fn refresh_mcp_status() -> Result<(), ServerFnError> {
    use chrono::Utc;
    use compliance_core::models::McpServerStatus;
    use mongodb::bson::doc;

    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let mut cursor = state
        .db
        .mcp_servers()
        .find(doc! {})
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(5))
        .build()
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    while cursor
        .advance()
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?
    {
        let server: compliance_core::models::McpServerConfig = cursor
            .deserialize_current()
            .map_err(|e| ServerFnError::new(e.to_string()))?;

        let Some(oid) = server.id else { continue };

        // Derive health URL from the endpoint (replace trailing /mcp with /health)
        let health_url = if server.endpoint_url.ends_with("/mcp") {
            format!(
                "{}health",
                &server.endpoint_url[..server.endpoint_url.len() - 3]
            )
        } else {
            format!("{}/health", server.endpoint_url.trim_end_matches('/'))
        };

        let new_status = match client.get(&health_url).send().await {
            Ok(resp) if resp.status().is_success() => McpServerStatus::Running,
            _ => McpServerStatus::Stopped,
        };

        let status_bson = match bson::to_bson(&new_status) {
            Ok(b) => b,
            Err(_) => continue,
        };

        let _ = state
            .db
            .mcp_servers()
            .update_one(
                doc! { "_id": oid },
                doc! { "$set": { "status": status_bson, "updated_at": Utc::now().to_rfc3339() } },
            )
            .await;
    }

    Ok(())
}

#[server]
pub async fn delete_mcp_server(server_id: String) -> Result<(), ServerFnError> {
    use mongodb::bson::doc;

    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let oid = bson::oid::ObjectId::parse_str(&server_id)
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    state
        .db
        .mcp_servers()
        .delete_one(doc! { "_id": oid })
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    Ok(())
}

#[server]
pub async fn regenerate_mcp_token(server_id: String) -> Result<String, ServerFnError> {
    use chrono::Utc;
    use mongodb::bson::doc;

    let state: super::server_state::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let oid = bson::oid::ObjectId::parse_str(&server_id)
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    let new_token = format!("mcp_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));

    state
        .db
        .mcp_servers()
        .update_one(
            doc! { "_id": oid },
            doc! { "$set": { "access_token": &new_token, "updated_at": Utc::now().to_rfc3339() } },
        )
        .await
        .map_err(|e| ServerFnError::new(e.to_string()))?;

    Ok(new_token)
}