compliance-scanner-agent/compliance-agent/src/api/handlers/mcp_tokens.rs

//! `/api/v1/mcp-tokens` — per-tenant API tokens for the MCP server.
//!
//! These are opaque static bearers issued via the dashboard (or a
//! direct curl with a KC JWT) and copied into LLM clients (Claude
//! Desktop / Cursor / ChatGPT). The MCP server hashes incoming bearers
//! and looks them up in the cross-tenant `<prefix>__admin.mcp_tokens`
//! collection to derive the tenant_id for routing.
//!
//! The raw token is shown to the caller exactly once at creation; the
//! database only ever stores the SHA-256 hash. Revocation is a soft
//! delete (sets `revoked: true`) so the audit log keeps the record.

use axum::extract::{Extension, Path};
use axum::http::StatusCode;
use axum::Json;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use compliance_core::models::{McpToken, McpTokenView};
use compliance_core::tenant_ctx::TenantCtx;
use mongodb::bson::doc;
use rand::RngCore;
use sha2::{Digest, Sha256};

use super::dto::{AgentExt, ApiResponse};

/// Mongo collection name inside the admin DB.
const COLLECTION: &str = "mcp_tokens";

/// Token prefix the MCP server expects on every bearer.
const TOKEN_PREFIX: &str = "mcpt_";

/// Bytes of randomness behind each token. 32 → ~256 bits.
/// Encoded as URL-safe base64 without padding → 43 chars.
/// Combined with `mcpt_` → 48-char tokens.
const TOKEN_RAND_BYTES: usize = 32;

#[derive(serde::Deserialize)]
pub struct CreateMcpTokenRequest {
    pub name: String,
}

/// Returned exactly once at creation. The `token` field is gone from
/// the listing endpoint — the user must save it now.
#[derive(serde::Serialize)]
pub struct CreateMcpTokenResponse {
    pub token: String,
    pub view: McpTokenView,
}

/// `POST /api/v1/mcp-tokens` — mint a new token for the caller's tenant.
#[tracing::instrument(skip_all)]
pub async fn create_mcp_token(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Json(req): Json<CreateMcpTokenRequest>,
) -> Result<Json<CreateMcpTokenResponse>, StatusCode> {
    if req.name.trim().is_empty() {
        return Err(StatusCode::BAD_REQUEST);
    }
    let raw = generate_token();
    let token_hash = sha256_hex(&raw);
    let token_prefix: String = raw.chars().take(12).collect();

    let mut token = McpToken {
        id: None,
        token_hash,
        token_prefix,
        tenant_id: tenant.0.tenant_id.clone(),
        name: req.name.trim().to_string(),
        created_by: tenant.0.user_id.clone(),
        created_at: chrono::Utc::now(),
        last_used_at: None,
        revoked: false,
    };

    let col = agent.db_pool.admin_db().collection::<McpToken>(COLLECTION);
    let res = col.insert_one(&token).await.map_err(|e| {
        tracing::error!("Failed to insert MCP token: {e}");
        StatusCode::INTERNAL_SERVER_ERROR
    })?;
    token.id = res.inserted_id.as_object_id();

    Ok(Json(CreateMcpTokenResponse {
        view: McpTokenView::from(&token),
        token: raw,
    }))
}

/// `GET /api/v1/mcp-tokens` — list tokens for the caller's tenant.
/// Hash is never returned; only metadata + the 12-char prefix so the
/// user can identify which row is which.
#[tracing::instrument(skip_all)]
pub async fn list_mcp_tokens(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
) -> Result<Json<ApiResponse<Vec<McpTokenView>>>, StatusCode> {
    let col = agent.db_pool.admin_db().collection::<McpToken>(COLLECTION);
    let mut cursor = col
        .find(doc! { "tenant_id": &tenant.0.tenant_id })
        .sort(doc! { "created_at": -1 })
        .await
        .map_err(|e| {
            tracing::error!("Failed to list MCP tokens: {e}");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;
    let mut out = Vec::new();
    while cursor.advance().await.map_err(|e| {
        tracing::warn!("MCP tokens cursor advance failed: {e}");
        StatusCode::INTERNAL_SERVER_ERROR
    })? {
        match cursor.deserialize_current() {
            Ok(t) => out.push(McpTokenView::from(&t)),
            Err(e) => tracing::warn!("Failed to deserialize MCP token: {e}"),
        }
    }
    Ok(Json(ApiResponse {
        data: out,
        total: None,
        page: None,
    }))
}

/// `DELETE /api/v1/mcp-tokens/{id}` — revoke (soft delete).
/// Scoped to the caller's tenant: a user can't revoke another tenant's
/// token even if they guess its id.
#[tracing::instrument(skip_all, fields(id = %id))]
pub async fn revoke_mcp_token(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, StatusCode> {
    let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
    let col = agent.db_pool.admin_db().collection::<McpToken>(COLLECTION);
    let result = col
        .update_one(
            doc! { "_id": oid, "tenant_id": &tenant.0.tenant_id },
            doc! { "$set": { "revoked": true } },
        )
        .await
        .map_err(|e| {
            tracing::error!("Failed to revoke MCP token: {e}");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;
    if result.matched_count == 0 {
        return Err(StatusCode::NOT_FOUND);
    }
    Ok(Json(serde_json::json!({ "status": "revoked" })))
}

/// 32 bytes random → URL-safe base64 → 43 chars, no padding.
/// Prefixed with `mcpt_` so the MCP server can sniff the format
/// before bothering with the DB lookup.
fn generate_token() -> String {
    let mut bytes = [0u8; TOKEN_RAND_BYTES];
    rand::rng().fill_bytes(&mut bytes);
    format!("{TOKEN_PREFIX}{}", URL_SAFE_NO_PAD.encode(bytes))
}

fn sha256_hex(s: &str) -> String {
    let mut h = Sha256::new();
    h.update(s.as_bytes());
    hex::encode(h.finalize())
}

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

    #[test]
    fn generated_tokens_are_unique_and_prefixed() {
        let a = generate_token();
        let b = generate_token();
        assert_ne!(a, b);
        assert!(a.starts_with(TOKEN_PREFIX));
        assert!(b.starts_with(TOKEN_PREFIX));
        // 5 + 43 = 48 chars
        assert_eq!(a.len(), 5 + 43);
    }

    #[test]
    fn sha256_is_stable_and_64_hex() {
        let h = sha256_hex("mcpt_abc");
        assert_eq!(h.len(), 64);
        assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
        assert_eq!(sha256_hex("mcpt_abc"), h);
    }
}