compliance-scanner-agent/compliance-mcp/src/tools/findings.rs

use mongodb::bson::doc;
use rmcp::{model::*, ErrorData as McpError};
use schemars::JsonSchema;
use serde::Deserialize;

use crate::database::Database;

const MAX_LIMIT: i64 = 200;
const DEFAULT_LIMIT: i64 = 50;

fn cap_limit(limit: Option<i64>) -> i64 {
    limit.unwrap_or(DEFAULT_LIMIT).clamp(1, MAX_LIMIT)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cap_limit_default() {
        assert_eq!(cap_limit(None), DEFAULT_LIMIT);
    }

    #[test]
    fn cap_limit_normal_value() {
        assert_eq!(cap_limit(Some(100)), 100);
    }

    #[test]
    fn cap_limit_exceeds_max() {
        assert_eq!(cap_limit(Some(500)), MAX_LIMIT);
        assert_eq!(cap_limit(Some(201)), MAX_LIMIT);
    }

    #[test]
    fn cap_limit_zero_clamped_to_one() {
        assert_eq!(cap_limit(Some(0)), 1);
    }

    #[test]
    fn cap_limit_negative_clamped_to_one() {
        assert_eq!(cap_limit(Some(-10)), 1);
    }

    #[test]
    fn cap_limit_boundary_values() {
        assert_eq!(cap_limit(Some(1)), 1);
        assert_eq!(cap_limit(Some(MAX_LIMIT)), MAX_LIMIT);
    }

    #[test]
    fn list_findings_params_deserialize() {
        let json = serde_json::json!({
            "repo_id": "abc",
            "severity": "high",
            "status": "open",
            "scan_type": "sast",
            "limit": 25
        });
        let params: ListFindingsParams = serde_json::from_value(json).unwrap();
        assert_eq!(params.repo_id.as_deref(), Some("abc"));
        assert_eq!(params.severity.as_deref(), Some("high"));
        assert_eq!(params.status.as_deref(), Some("open"));
        assert_eq!(params.scan_type.as_deref(), Some("sast"));
        assert_eq!(params.limit, Some(25));
    }

    #[test]
    fn list_findings_params_all_optional() {
        let json = serde_json::json!({});
        let params: ListFindingsParams = serde_json::from_value(json).unwrap();
        assert!(params.repo_id.is_none());
        assert!(params.severity.is_none());
        assert!(params.status.is_none());
        assert!(params.scan_type.is_none());
        assert!(params.limit.is_none());
    }

    #[test]
    fn get_finding_params_deserialize() {
        let json = serde_json::json!({ "id": "507f1f77bcf86cd799439011" });
        let params: GetFindingParams = serde_json::from_value(json).unwrap();
        assert_eq!(params.id, "507f1f77bcf86cd799439011");
    }

    #[test]
    fn findings_summary_params_deserialize() {
        let json = serde_json::json!({ "repo_id": "r1" });
        let params: FindingsSummaryParams = serde_json::from_value(json).unwrap();
        assert_eq!(params.repo_id.as_deref(), Some("r1"));

        let json2 = serde_json::json!({});
        let params2: FindingsSummaryParams = serde_json::from_value(json2).unwrap();
        assert!(params2.repo_id.is_none());
    }
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct ListFindingsParams {
    /// Filter by repository ID
    pub repo_id: Option<String>,
    /// Filter by severity: info, low, medium, high, critical
    pub severity: Option<String>,
    /// Filter by status: open, triaged, false_positive, resolved, ignored
    pub status: Option<String>,
    /// Filter by scan type: sast, sbom, cve, gdpr, oauth, secret_detection, lint, code_review
    pub scan_type: Option<String>,
    /// Maximum number of results (default 50, max 200)
    pub limit: Option<i64>,
}

pub async fn list_findings(
    db: &Database,
    params: ListFindingsParams,
) -> Result<CallToolResult, McpError> {
    let mut filter = doc! {};
    if let Some(ref repo_id) = params.repo_id {
        filter.insert("repo_id", repo_id);
    }
    if let Some(ref severity) = params.severity {
        filter.insert("severity", severity);
    }
    if let Some(ref status) = params.status {
        filter.insert("status", status);
    }
    if let Some(ref scan_type) = params.scan_type {
        filter.insert("scan_type", scan_type);
    }

    let limit = cap_limit(params.limit);

    let mut cursor = db
        .findings()
        .find(filter)
        .sort(doc! { "created_at": -1 })
        .limit(limit)
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;

    let mut results = Vec::new();
    while cursor
        .advance()
        .await
        .map_err(|e| McpError::internal_error(format!("cursor error: {e}"), None))?
    {
        let finding = cursor
            .deserialize_current()
            .map_err(|e| McpError::internal_error(format!("deserialize error: {e}"), None))?;
        results.push(finding);
    }

    let json = serde_json::to_string_pretty(&results)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct GetFindingParams {
    /// Finding ID (MongoDB ObjectId hex string)
    pub id: String,
}

pub async fn get_finding(
    db: &Database,
    params: GetFindingParams,
) -> Result<CallToolResult, McpError> {
    let oid = bson::oid::ObjectId::parse_str(&params.id)
        .map_err(|e| McpError::invalid_params(format!("invalid ObjectId: {e}"), None))?;

    let finding = db
        .findings()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?
        .ok_or_else(|| McpError::invalid_params("finding not found", None))?;

    let json = serde_json::to_string_pretty(&finding)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct FindingsSummaryParams {
    /// Filter by repository ID
    pub repo_id: Option<String>,
}

#[derive(serde::Serialize)]
struct SeverityCount {
    severity: String,
    count: u64,
}

pub async fn findings_summary(
    db: &Database,
    params: FindingsSummaryParams,
) -> Result<CallToolResult, McpError> {
    let mut base_filter = doc! {};
    if let Some(ref repo_id) = params.repo_id {
        base_filter.insert("repo_id", repo_id);
    }

    let severities = ["critical", "high", "medium", "low", "info"];
    let mut counts = Vec::new();

    for sev in &severities {
        let mut filter = base_filter.clone();
        filter.insert("severity", sev);
        let count = db
            .findings()
            .count_documents(filter)
            .await
            .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;
        counts.push(SeverityCount {
            severity: sev.to_string(),
            count,
        });
    }

    let total: u64 = counts.iter().map(|c| c.count).sum();

    let mut status_counts = Vec::new();
    for status in &["open", "triaged", "false_positive", "resolved", "ignored"] {
        let mut filter = base_filter.clone();
        filter.insert("status", status);
        let count = db
            .findings()
            .count_documents(filter)
            .await
            .map_err(|e| McpError::internal_error(format!("DB error: {e}"), None))?;
        status_counts.push(serde_json::json!({ "status": status, "count": count }));
    }

    let summary = serde_json::json!({
        "total": total,
        "by_severity": counts,
        "by_status": status_counts,
    });

    let json = serde_json::to_string_pretty(&summary)
        .map_err(|e| McpError::internal_error(format!("json error: {e}"), None))?;

    Ok(CallToolResult::success(vec![Content::text(json)]))
}