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

use axum::routing::{get, post};
use axum::{middleware, Extension};
use dioxus::prelude::*;
use time::Duration;
use tower_sessions::{cookie::Key, MemoryStore, SessionManagerLayer};

use compliance_core::models::{McpServerConfig, McpServerStatus, McpTransport};
use mongodb::bson::doc;

use super::config;
use super::database::Database;
use super::error::DashboardError;
use super::keycloak_config::KeycloakConfig;
use super::server_state::{ServerState, ServerStateInner};
use super::{auth_callback, auth_login, logout, require_auth, PendingOAuthStore};

pub fn server_start(app: fn() -> Element) -> Result<(), DashboardError> {
    tokio::runtime::Runtime::new()
        .map_err(|e| DashboardError::Other(e.to_string()))?
        .block_on(async move {
            dotenvy::dotenv().ok();

            let config = config::load_config()?;
            let keycloak: Option<&'static KeycloakConfig> =
                KeycloakConfig::from_env().map(|kc| &*Box::leak(Box::new(kc)));
            let db = Database::connect(&config.mongodb_uri, &config.mongodb_database).await?;

            // Seed default MCP server configs
            seed_default_mcp_servers(&db, config.mcp_endpoint_url.as_deref()).await;

            if let Some(kc) = keycloak {
                tracing::info!("Keycloak configured for realm '{}'", kc.realm);
            } else {
                tracing::warn!("Keycloak not configured - dashboard is unprotected");
            }

            let server_state: ServerState = ServerStateInner {
                agent_api_url: config.agent_api_url.clone(),
                db,
                config,
                keycloak,
            }
            .into();

            // Session layer
            let key = Key::generate();
            let store = MemoryStore::default();
            let session = SessionManagerLayer::new(store)
                .with_secure(false)
                .with_same_site(tower_sessions::cookie::SameSite::Lax)
                .with_expiry(tower_sessions::Expiry::OnInactivity(Duration::hours(24)))
                .with_signed(key);

            let port = dioxus_cli_config::server_port().unwrap_or(8080);
            let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
            let listener = tokio::net::TcpListener::bind(addr)
                .await
                .map_err(|e| DashboardError::Other(format!("Failed to bind: {e}")))?;

            tracing::info!("Dashboard server listening on {addr}");

            let router = axum::Router::new()
                .route("/auth", get(auth_login))
                .route("/auth/callback", get(auth_callback))
                .route("/logout", get(logout))
                // Webhook proxy: forward to agent (no auth required)
                .route("/webhook/{platform}/{repo_id}", post(webhook_proxy))
                .serve_dioxus_application(ServeConfig::new(), app)
                .layer(Extension(PendingOAuthStore::default()))
                .layer(middleware::from_fn(require_auth))
                .layer(Extension(server_state))
                .layer(session);

            axum::serve(listener, router.into_make_service())
                .await
                .map_err(|e| DashboardError::Other(format!("Server error: {e}")))?;

            Ok(())
        })
}

/// Forward incoming webhooks to the agent's webhook server.
/// The dashboard acts as a public-facing proxy so the agent isn't exposed directly.
async fn webhook_proxy(
    Extension(state): Extension<ServerState>,
    axum::extract::Path((platform, repo_id)): axum::extract::Path<(String, String)>,
    headers: axum::http::HeaderMap,
    body: axum::body::Bytes,
) -> axum::http::StatusCode {
    // The agent_api_url typically looks like "http://agent:3001" or "http://localhost:3001"
    // Webhook routes are on the same server, so strip any trailing path
    let base = state.agent_api_url.trim_end_matches('/');
    // Remove /api/v1 suffix if present to get base URL
    let base = base
        .strip_suffix("/api/v1")
        .or_else(|| base.strip_suffix("/api"))
        .unwrap_or(base);
    let agent_url = format!("{base}/webhook/{platform}/{repo_id}");

    // Forward all relevant headers
    let client = reqwest::Client::new();
    let mut req = client.post(&agent_url).body(body.to_vec());

    for (name, value) in &headers {
        let name_str = name.as_str().to_lowercase();
        // Forward platform-specific headers
        if name_str.starts_with("x-gitea-")
            || name_str.starts_with("x-github-")
            || name_str.starts_with("x-hub-")
            || name_str.starts_with("x-gitlab-")
            || name_str == "content-type"
        {
            if let Ok(v) = value.to_str() {
                req = req.header(name.as_str(), v);
            }
        }
    }

    match req.send().await {
        Ok(resp) => axum::http::StatusCode::from_u16(resp.status().as_u16())
            .unwrap_or(axum::http::StatusCode::BAD_GATEWAY),
        Err(e) => {
            tracing::error!("Webhook proxy failed: {e}");
            axum::http::StatusCode::BAD_GATEWAY
        }
    }
}

/// Seed three default MCP server configs (Findings, SBOM, DAST) if they don't already exist.
async fn seed_default_mcp_servers(db: &Database, mcp_endpoint_url: Option<&str>) {
    let endpoint = mcp_endpoint_url.unwrap_or("http://localhost:8090");

    let defaults = [
        (
            "Findings MCP",
            "Exposes security findings, triage data, and finding summaries to LLM agents",
            vec!["list_findings", "get_finding", "findings_summary"],
        ),
        (
            "SBOM MCP",
            "Exposes software bill of materials and vulnerability reports to LLM agents",
            vec!["list_sbom_packages", "sbom_vuln_report"],
        ),
        (
            "DAST MCP",
            "Exposes DAST scan findings and scan summaries to LLM agents",
            vec!["list_dast_findings", "dast_scan_summary"],
        ),
    ];

    let collection = db.mcp_servers();

    let expected_url = format!("{endpoint}/mcp");

    for (name, description, tools) in defaults {
        // If it already exists, update the endpoint URL if it changed
        if let Ok(Some(existing)) = collection.find_one(doc! { "name": name }).await {
            if existing.endpoint_url != expected_url {
                let _ = collection
                    .update_one(
                        doc! { "name": name },
                        doc! { "$set": { "endpoint_url": &expected_url } },
                    )
                    .await;
                tracing::info!(
                    "Updated MCP server '{name}' endpoint: {} -> {expected_url}",
                    existing.endpoint_url
                );
            }
            continue;
        }

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

        let server = McpServerConfig {
            id: None,
            name: name.to_string(),
            endpoint_url: format!("{endpoint}/mcp"),
            transport: McpTransport::Http,
            port: Some(8090),
            status: McpServerStatus::Stopped,
            access_token: token,
            tools_enabled: tools.into_iter().map(|s| s.to_string()).collect(),
            description: Some(description.to_string()),
            mongodb_uri: None,
            mongodb_database: None,
            created_at: now,
            updated_at: now,
        };

        match collection.insert_one(server).await {
            Ok(_) => tracing::info!("Seeded default MCP server: {name}"),
            Err(e) => tracing::warn!("Failed to seed MCP server '{name}': {e}"),
        }
    }
}