secrets/crates/secrets-mcp/src/tools.rs

use std::sync::Arc;

use anyhow::Result;
use rmcp::{
    RoleServer, ServerHandler,
    handler::server::wrapper::Parameters,
    model::{
        CallToolResult, Content, Implementation, InitializeResult, ProtocolVersion,
        ServerCapabilities,
    },
    service::RequestContext,
    tool, tool_handler, tool_router,
};
use schemars::JsonSchema;
use serde::Deserialize;
use sqlx::PgPool;
use uuid::Uuid;

use secrets_core::service::{
    add::{AddParams, run as svc_add},
    delete::{DeleteParams, run as svc_delete},
    export::{ExportParams, export as svc_export},
    get_secret::{get_all_secrets, get_secret_field},
    history::run as svc_history,
    rollback::run as svc_rollback,
    search::{SearchParams, run as svc_search},
    update::{UpdateParams, run as svc_update},
};

use crate::auth::AuthUser;

// ── Shared state ──────────────────────────────────────────────────────────────

#[derive(Clone)]
pub struct SecretsService {
    pub pool: Arc<PgPool>,
    pub tool_router: rmcp::handler::server::router::tool::ToolRouter<SecretsService>,
}

impl SecretsService {
    pub fn new(pool: Arc<PgPool>) -> Self {
        Self {
            pool,
            tool_router: Self::tool_router(),
        }
    }

    /// Extract user_id from the HTTP request parts injected by auth middleware.
    fn user_id_from_ctx(ctx: &RequestContext<RoleServer>) -> Result<Option<Uuid>, rmcp::ErrorData> {
        let parts = ctx
            .extensions
            .get::<http::request::Parts>()
            .ok_or_else(|| rmcp::ErrorData::internal_error("Missing HTTP parts", None))?;
        Ok(parts.extensions.get::<AuthUser>().map(|a| a.user_id))
    }

    /// Get the authenticated user_id (returns error if not authenticated).
    fn require_user_id(ctx: &RequestContext<RoleServer>) -> Result<Uuid, rmcp::ErrorData> {
        let parts = ctx
            .extensions
            .get::<http::request::Parts>()
            .ok_or_else(|| rmcp::ErrorData::internal_error("Missing HTTP parts", None))?;
        parts
            .extensions
            .get::<AuthUser>()
            .map(|a| a.user_id)
            .ok_or_else(|| rmcp::ErrorData::invalid_request("Unauthorized: API key required", None))
    }

    /// Extract the 32-byte encryption key from the X-Encryption-Key request header.
    /// The header value must be 64 lowercase hex characters (PBKDF2-derived key).
    fn extract_enc_key(ctx: &RequestContext<RoleServer>) -> Result<[u8; 32], rmcp::ErrorData> {
        let parts = ctx
            .extensions
            .get::<http::request::Parts>()
            .ok_or_else(|| rmcp::ErrorData::internal_error("Missing HTTP parts", None))?;
        let hex_str = parts
            .headers
            .get("x-encryption-key")
            .ok_or_else(|| {
                rmcp::ErrorData::invalid_request(
                    "Missing X-Encryption-Key header. \
                     Set this to your 64-char hex encryption key derived from your passphrase.",
                    None,
                )
            })?
            .to_str()
            .map_err(|_| {
                rmcp::ErrorData::invalid_request("Invalid X-Encryption-Key header value", None)
            })?;
        secrets_core::crypto::extract_key_from_hex(hex_str)
            .map_err(|e| rmcp::ErrorData::invalid_request(e.to_string(), None))
    }

    /// Require both user_id and encryption key.
    fn require_user_and_key(
        ctx: &RequestContext<RoleServer>,
    ) -> Result<(Uuid, [u8; 32]), rmcp::ErrorData> {
        let user_id = Self::require_user_id(ctx)?;
        let key = Self::extract_enc_key(ctx)?;
        Ok((user_id, key))
    }
}

// ── Tool parameter types ──────────────────────────────────────────────────────

#[derive(Debug, Deserialize, JsonSchema)]
struct SearchInput {
    #[schemars(description = "Namespace filter (e.g. 'refining', 'ricnsmart')")]
    namespace: Option<String>,
    #[schemars(description = "Kind filter (e.g. 'server', 'service', 'key')")]
    kind: Option<String>,
    #[schemars(description = "Exact record name")]
    name: Option<String>,
    #[schemars(description = "Tag filters (all must match)")]
    tags: Option<Vec<String>>,
    #[schemars(description = "Fuzzy search across name, namespace, kind, tags, metadata")]
    query: Option<String>,
    #[schemars(description = "Return only summary fields (name/tags/desc/updated_at)")]
    summary: Option<bool>,
    #[schemars(description = "Sort order: 'name' (default), 'updated', 'created'")]
    sort: Option<String>,
    #[schemars(description = "Max results (default 20)")]
    limit: Option<u32>,
    #[schemars(description = "Pagination offset (default 0)")]
    offset: Option<u32>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct GetSecretInput {
    #[schemars(description = "Namespace of the entry")]
    namespace: String,
    #[schemars(description = "Kind of the entry (e.g. 'server', 'service')")]
    kind: String,
    #[schemars(description = "Name of the entry")]
    name: String,
    #[schemars(description = "Specific field to retrieve. If omitted, returns all fields.")]
    field: Option<String>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct AddInput {
    #[schemars(description = "Namespace")]
    namespace: String,
    #[schemars(description = "Kind (e.g. 'server', 'service', 'key')")]
    kind: String,
    #[schemars(description = "Unique name within namespace+kind")]
    name: String,
    #[schemars(description = "Tags for this entry")]
    tags: Option<Vec<String>>,
    #[schemars(description = "Metadata fields as 'key=value' or 'key:=json' strings")]
    meta: Option<Vec<String>>,
    #[schemars(description = "Secret fields as 'key=value' strings")]
    secrets: Option<Vec<String>>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct UpdateInput {
    #[schemars(description = "Namespace")]
    namespace: String,
    #[schemars(description = "Kind")]
    kind: String,
    #[schemars(description = "Name")]
    name: String,
    #[schemars(description = "Tags to add")]
    add_tags: Option<Vec<String>>,
    #[schemars(description = "Tags to remove")]
    remove_tags: Option<Vec<String>>,
    #[schemars(description = "Metadata fields to update/add as 'key=value' strings")]
    meta: Option<Vec<String>>,
    #[schemars(description = "Metadata field keys to remove")]
    remove_meta: Option<Vec<String>>,
    #[schemars(description = "Secret fields to update/add as 'key=value' strings")]
    secrets: Option<Vec<String>>,
    #[schemars(description = "Secret field keys to remove")]
    remove_secrets: Option<Vec<String>>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct DeleteInput {
    #[schemars(description = "Namespace")]
    namespace: String,
    #[schemars(description = "Kind filter (required for single delete)")]
    kind: Option<String>,
    #[schemars(description = "Exact name to delete. Omit for bulk delete by namespace+kind.")]
    name: Option<String>,
    #[schemars(description = "Preview deletions without writing")]
    dry_run: Option<bool>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct HistoryInput {
    #[schemars(description = "Namespace")]
    namespace: String,
    #[schemars(description = "Kind")]
    kind: String,
    #[schemars(description = "Name")]
    name: String,
    #[schemars(description = "Max history entries to return (default 20)")]
    limit: Option<u32>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct RollbackInput {
    #[schemars(description = "Namespace")]
    namespace: String,
    #[schemars(description = "Kind")]
    kind: String,
    #[schemars(description = "Name")]
    name: String,
    #[schemars(description = "Target version number. Omit to restore the most recent snapshot.")]
    to_version: Option<i64>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct ExportInput {
    #[schemars(description = "Namespace filter")]
    namespace: Option<String>,
    #[schemars(description = "Kind filter")]
    kind: Option<String>,
    #[schemars(description = "Exact name filter")]
    name: Option<String>,
    #[schemars(description = "Tag filters")]
    tags: Option<Vec<String>>,
    #[schemars(description = "Fuzzy query")]
    query: Option<String>,
    #[schemars(description = "Export format: 'json' (default), 'toml', 'yaml'")]
    format: Option<String>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct EnvMapInput {
    #[schemars(description = "Namespace filter")]
    namespace: Option<String>,
    #[schemars(description = "Kind filter")]
    kind: Option<String>,
    #[schemars(description = "Exact name filter")]
    name: Option<String>,
    #[schemars(description = "Tag filters")]
    tags: Option<Vec<String>>,
    #[schemars(description = "Only include these secret fields")]
    only_fields: Option<Vec<String>>,
    #[schemars(description = "Environment variable name prefix")]
    prefix: Option<String>,
}

// ── Tool implementations ──────────────────────────────────────────────────────

#[tool_router]
impl SecretsService {
    #[tool(
        description = "Search entries in the secrets store. Returns entries with metadata and \
        secret field names (not values). Use secrets_get to decrypt secret values."
    )]
    async fn secrets_search(
        &self,
        Parameters(input): Parameters<SearchInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let user_id = Self::user_id_from_ctx(&ctx)?;
        let tags = input.tags.unwrap_or_default();
        let result = svc_search(
            &self.pool,
            SearchParams {
                namespace: input.namespace.as_deref(),
                kind: input.kind.as_deref(),
                name: input.name.as_deref(),
                tags: &tags,
                query: input.query.as_deref(),
                sort: input.sort.as_deref().unwrap_or("name"),
                limit: input.limit.unwrap_or(20),
                offset: input.offset.unwrap_or(0),
                user_id,
            },
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let summary = input.summary.unwrap_or(false);
        let entries: Vec<serde_json::Value> = result
            .entries
            .iter()
            .map(|e| {
                if summary {
                    serde_json::json!({
                        "namespace": e.namespace,
                        "kind": e.kind,
                        "name": e.name,
                        "tags": e.tags,
                        "desc": e.metadata.get("desc").or_else(|| e.metadata.get("url"))
                            .and_then(|v| v.as_str()).unwrap_or(""),
                        "updated_at": e.updated_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                    })
                } else {
                    let schema: Vec<&str> = result
                        .secret_schemas
                        .get(&e.id)
                        .map(|f| f.iter().map(|s| s.field_name.as_str()).collect())
                        .unwrap_or_default();
                    serde_json::json!({
                        "id": e.id,
                        "namespace": e.namespace,
                        "kind": e.kind,
                        "name": e.name,
                        "tags": e.tags,
                        "metadata": e.metadata,
                        "secret_fields": schema,
                        "version": e.version,
                        "updated_at": e.updated_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                    })
                }
            })
            .collect();

        let json = serde_json::to_string_pretty(&entries).unwrap_or_else(|_| "[]".to_string());
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "Get decrypted secret field values for an entry. Requires your \
        encryption key via X-Encryption-Key header (64 hex chars, PBKDF2-derived). \
        Returns all fields, or a specific field if 'field' is provided."
    )]
    async fn secrets_get(
        &self,
        Parameters(input): Parameters<GetSecretInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;

        if let Some(field_name) = &input.field {
            let value = get_secret_field(
                &self.pool,
                &input.namespace,
                &input.kind,
                &input.name,
                field_name,
                &user_key,
                Some(user_id),
            )
            .await
            .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

            let result = serde_json::json!({ field_name: value });
            let json = serde_json::to_string_pretty(&result).unwrap_or_default();
            Ok(CallToolResult::success(vec![Content::text(json)]))
        } else {
            let secrets = get_all_secrets(
                &self.pool,
                &input.namespace,
                &input.kind,
                &input.name,
                &user_key,
                Some(user_id),
            )
            .await
            .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

            let json = serde_json::to_string_pretty(&secrets).unwrap_or_default();
            Ok(CallToolResult::success(vec![Content::text(json)]))
        }
    }

    #[tool(
        description = "Add or upsert an entry with metadata and encrypted secret fields. \
        Requires X-Encryption-Key header. \
        Meta and secret values use 'key=value', 'key=@file', or 'key:=<json>' format."
    )]
    async fn secrets_add(
        &self,
        Parameters(input): Parameters<AddInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;

        let tags = input.tags.unwrap_or_default();
        let meta = input.meta.unwrap_or_default();
        let secrets = input.secrets.unwrap_or_default();

        let result = svc_add(
            &self.pool,
            AddParams {
                namespace: &input.namespace,
                kind: &input.kind,
                name: &input.name,
                tags: &tags,
                meta_entries: &meta,
                secret_entries: &secrets,
                user_id: Some(user_id),
            },
            &user_key,
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&result).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "Incrementally update an existing entry. Requires X-Encryption-Key header. \
        Only the fields you specify are changed; everything else is preserved."
    )]
    async fn secrets_update(
        &self,
        Parameters(input): Parameters<UpdateInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;

        let add_tags = input.add_tags.unwrap_or_default();
        let remove_tags = input.remove_tags.unwrap_or_default();
        let meta = input.meta.unwrap_or_default();
        let remove_meta = input.remove_meta.unwrap_or_default();
        let secrets = input.secrets.unwrap_or_default();
        let remove_secrets = input.remove_secrets.unwrap_or_default();

        let result = svc_update(
            &self.pool,
            UpdateParams {
                namespace: &input.namespace,
                kind: &input.kind,
                name: &input.name,
                add_tags: &add_tags,
                remove_tags: &remove_tags,
                meta_entries: &meta,
                remove_meta: &remove_meta,
                secret_entries: &secrets,
                remove_secrets: &remove_secrets,
                user_id: Some(user_id),
            },
            &user_key,
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&result).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "Delete one entry (specify namespace+kind+name) or bulk delete all \
        entries matching namespace+kind. Use dry_run=true to preview."
    )]
    async fn secrets_delete(
        &self,
        Parameters(input): Parameters<DeleteInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let user_id = Self::user_id_from_ctx(&ctx)?;

        let result = svc_delete(
            &self.pool,
            DeleteParams {
                namespace: &input.namespace,
                kind: input.kind.as_deref(),
                name: input.name.as_deref(),
                dry_run: input.dry_run.unwrap_or(false),
                user_id,
            },
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&result).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "View change history for an entry. Returns a list of versions with \
        actions and timestamps."
    )]
    async fn secrets_history(
        &self,
        Parameters(input): Parameters<HistoryInput>,
        _ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let result = svc_history(
            &self.pool,
            &input.namespace,
            &input.kind,
            &input.name,
            input.limit.unwrap_or(20),
            None,
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&result).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "Rollback an entry to a previous version. Requires X-Encryption-Key header. \
        Omit to_version to restore the most recent snapshot."
    )]
    async fn secrets_rollback(
        &self,
        Parameters(input): Parameters<RollbackInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;

        let result = svc_rollback(
            &self.pool,
            &input.namespace,
            &input.kind,
            &input.name,
            input.to_version,
            &user_key,
            Some(user_id),
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&result).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(
        description = "Export matching entries with decrypted secrets as JSON/TOML/YAML string. \
        Requires X-Encryption-Key header. Useful for backup or data migration."
    )]
    async fn secrets_export(
        &self,
        Parameters(input): Parameters<ExportInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;
        let tags = input.tags.unwrap_or_default();
        let format = input.format.as_deref().unwrap_or("json");

        let data = svc_export(
            &self.pool,
            ExportParams {
                namespace: input.namespace.as_deref(),
                kind: input.kind.as_deref(),
                name: input.name.as_deref(),
                tags: &tags,
                query: input.query.as_deref(),
                no_secrets: false,
                user_id: Some(user_id),
            },
            Some(&user_key),
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let serialized = format
            .parse::<secrets_core::models::ExportFormat>()
            .and_then(|fmt| fmt.serialize(&data))
            .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

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

    #[tool(
        description = "Build the environment variable map from entry secrets with decrypted \
        plaintext values. Requires X-Encryption-Key header. \
        Returns a JSON object of VAR_NAME -> plaintext_value ready for injection."
    )]
    async fn secrets_env_map(
        &self,
        Parameters(input): Parameters<EnvMapInput>,
        ctx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, rmcp::ErrorData> {
        let (user_id, user_key) = Self::require_user_and_key(&ctx)?;
        let tags = input.tags.unwrap_or_default();
        let only_fields = input.only_fields.unwrap_or_default();

        let env_map = secrets_core::service::env_map::build_env_map(
            &self.pool,
            input.namespace.as_deref(),
            input.kind.as_deref(),
            input.name.as_deref(),
            &tags,
            &only_fields,
            input.prefix.as_deref().unwrap_or(""),
            &user_key,
            Some(user_id),
        )
        .await
        .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?;

        let json = serde_json::to_string_pretty(&env_map).unwrap_or_default();
        Ok(CallToolResult::success(vec![Content::text(json)]))
    }
}

// ── ServerHandler ─────────────────────────────────────────────────────────────

#[tool_handler]
impl ServerHandler for SecretsService {
    fn get_info(&self) -> InitializeResult {
        let mut info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build());
        info.server_info = Implementation::new("secrets-mcp", env!("CARGO_PKG_VERSION"));
        info.protocol_version = ProtocolVersion::V_2025_03_26;
        info.instructions = Some(
            "Manage cross-device secrets and configuration securely. \
             Data is encrypted with your passphrase-derived key. \
             Include your 64-char hex key in the X-Encryption-Key header for all read/write operations. \
             Use secrets_search to discover entries (no key needed), \
             secrets_get to decrypt secret values, \
             and secrets_add/secrets_update to write encrypted secrets."
                .to_string(),
        );
        info
    }
}