compliance-scanner-agent/compliance-mcp/src/server.rs

use rmcp::{
    handler::server::wrapper::Parameters, model::*, tool, tool_handler, tool_router, ServerHandler,
};

use crate::auth::current_tenant_id;
use crate::database::{Database, DatabasePool};
use crate::tools::{dast, findings, pentest, sbom};

pub struct ComplianceMcpServer {
    pool: DatabasePool,
    #[allow(dead_code)]
    tool_router: rmcp::handler::server::router::tool::ToolRouter<Self>,
}

impl ComplianceMcpServer {
    /// Resolve the per-tenant `Database` from the bearer-set
    /// `task_local`. Every tool handler calls this; missing context
    /// surfaces as `internal_error` because it means the auth
    /// middleware was misconfigured (handler ran without scope).
    fn tenant_db(&self) -> Result<Database, rmcp::ErrorData> {
        let tenant_id = current_tenant_id().ok_or_else(|| {
            rmcp::ErrorData::internal_error(
                "no tenant context — bearer middleware not in chain".to_string(),
                None,
            )
        })?;
        Ok(self.pool.for_tenant_id(&tenant_id))
    }
}

#[tool_router]
impl ComplianceMcpServer {
    pub fn new(pool: DatabasePool) -> Self {
        Self {
            pool,
            tool_router: Self::tool_router(),
        }
    }

    // ── Findings ──────────────────────────────────────────

    #[tool(
        description = "List security findings with optional filters for repo, severity, status, and scan type"
    )]
    async fn list_findings(
        &self,
        Parameters(params): Parameters<findings::ListFindingsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::list_findings(&db, params).await
    }

    #[tool(description = "Get a single finding by its ID")]
    async fn get_finding(
        &self,
        Parameters(params): Parameters<findings::GetFindingParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::get_finding(&db, params).await
    }

    #[tool(description = "Get a summary of findings counts grouped by severity and status")]
    async fn findings_summary(
        &self,
        Parameters(params): Parameters<findings::FindingsSummaryParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        findings::findings_summary(&db, params).await
    }

    // ── SBOM ──────────────────────────────────────────────

    #[tool(
        description = "List SBOM packages with optional filters for repo, vulnerabilities, package manager, and license"
    )]
    async fn list_sbom_packages(
        &self,
        Parameters(params): Parameters<sbom::ListSbomPackagesParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        sbom::list_sbom_packages(&db, params).await
    }

    #[tool(
        description = "Generate a vulnerability report for a repository showing all packages with known CVEs"
    )]
    async fn sbom_vuln_report(
        &self,
        Parameters(params): Parameters<sbom::SbomVulnReportParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        sbom::sbom_vuln_report(&db, params).await
    }

    // ── DAST ──────────────────────────────────────────────

    #[tool(
        description = "List DAST findings with optional filters for target, scan run, severity, exploitability, and vulnerability type"
    )]
    async fn list_dast_findings(
        &self,
        Parameters(params): Parameters<dast::ListDastFindingsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        dast::list_dast_findings(&db, params).await
    }

    #[tool(description = "Get a summary of recent DAST scan runs and finding counts")]
    async fn dast_scan_summary(
        &self,
        Parameters(params): Parameters<dast::DastScanSummaryParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        dast::dast_scan_summary(&db, params).await
    }

    // ── Pentest ─────────────────────────────────────────────

    #[tool(
        description = "List AI pentest sessions with optional filters for target, status, and strategy"
    )]
    async fn list_pentest_sessions(
        &self,
        Parameters(params): Parameters<pentest::ListPentestSessionsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::list_pentest_sessions(&db, params).await
    }

    #[tool(description = "Get a single AI pentest session by its ID")]
    async fn get_pentest_session(
        &self,
        Parameters(params): Parameters<pentest::GetPentestSessionParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_pentest_session(&db, params).await
    }

    #[tool(
        description = "Get the attack chain DAG for a pentest session showing each tool invocation, its reasoning, and results"
    )]
    async fn get_attack_chain(
        &self,
        Parameters(params): Parameters<pentest::GetAttackChainParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_attack_chain(&db, params).await
    }

    #[tool(description = "Get chat messages from a pentest session")]
    async fn get_pentest_messages(
        &self,
        Parameters(params): Parameters<pentest::GetPentestMessagesParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::get_pentest_messages(&db, params).await
    }

    #[tool(
        description = "Get aggregated pentest statistics including running sessions, vulnerability counts, and severity distribution"
    )]
    async fn pentest_stats(
        &self,
        Parameters(params): Parameters<pentest::PentestStatsParams>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let db = self.tenant_db()?;
        pentest::pentest_stats(&db, params).await
    }
}

#[tool_handler]
impl ServerHandler for ComplianceMcpServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2024_11_05,
            capabilities: ServerCapabilities::builder()
                .enable_tools()
                .build(),
            server_info: Implementation::from_build_env(),
            instructions: Some(
                "Compliance Scanner MCP server. Query security findings, SBOM data, DAST results, and AI pentest sessions for your tenant."
                    .to_string(),
            ),
        }
    }
}