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

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

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

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

#[tool_router]
impl ComplianceMcpServer {
    pub fn new(db: Database) -> Self {
        Self {
            db,
            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> {
        findings::list_findings(&self.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> {
        findings::get_finding(&self.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> {
        findings::findings_summary(&self.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> {
        sbom::list_sbom_packages(&self.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> {
        sbom::sbom_vuln_report(&self.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> {
        dast::list_dast_findings(&self.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> {
        dast::dast_scan_summary(&self.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> {
        pentest::list_pentest_sessions(&self.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> {
        pentest::get_pentest_session(&self.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> {
        pentest::get_attack_chain(&self.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> {
        pentest::get_pentest_messages(&self.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> {
        pentest::pentest_stats(&self.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."
                    .to_string(),
            ),
        }
    }
}