symbi-runtime 1.21.1

//! MCP Client with Schema Verification
//!
//! Provides a secure MCP client that verifies tool schemas using SchemaPin

use async_trait::async_trait;
use sha2::Digest;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::time::{timeout, Duration};

use super::types::{
    McpClientConfig, McpClientError, McpTool, ToolDiscoveryEvent, ToolProvider,
    ToolVerificationRequest, ToolVerificationResponse, VerificationStatus,
};
use crate::integrations::schemapin::{
    LocalKeyStore, NativeSchemaPinClient, PinnedKey, SchemaPinClient,
};
use crate::integrations::tool_invocation::{
    DefaultToolInvocationEnforcer, InvocationContext, InvocationResult, ToolInvocationEnforcer,
    ToolInvocationError,
};

/// Trait for MCP client operations
#[async_trait]
pub trait McpClient: Send + Sync {
    /// Discover and verify a new tool
    async fn discover_tool(&self, tool: McpTool) -> Result<ToolDiscoveryEvent, McpClientError>;

    /// Get a verified tool by name
    async fn get_tool(&self, name: &str) -> Result<McpTool, McpClientError>;

    /// List all available tools
    async fn list_tools(&self) -> Result<Vec<McpTool>, McpClientError>;

    /// List only verified tools
    async fn list_verified_tools(&self) -> Result<Vec<McpTool>, McpClientError>;

    /// Verify a tool's schema
    async fn verify_tool(
        &self,
        request: ToolVerificationRequest,
    ) -> Result<ToolVerificationResponse, McpClientError>;

    /// Remove a tool from the client
    async fn remove_tool(&self, name: &str) -> Result<Option<McpTool>, McpClientError>;

    /// Execute a tool with verification enforcement
    async fn invoke_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
        context: InvocationContext,
    ) -> Result<InvocationResult, McpClientError>;
}

/// MCP client implementation with schema verification
pub struct SecureMcpClient {
    /// Client configuration
    config: McpClientConfig,
    /// SchemaPin client for verification
    schema_pin: Arc<dyn SchemaPinClient>,
    /// Local key store for TOFU
    key_store: Arc<LocalKeyStore>,
    /// Available tools (name -> tool)
    tools: Arc<RwLock<HashMap<String, McpTool>>>,
    /// Tool invocation enforcer
    enforcer: Arc<dyn ToolInvocationEnforcer>,
}

impl SecureMcpClient {
    /// Create a new secure MCP client
    pub fn new(
        config: McpClientConfig,
        schema_pin: Arc<dyn SchemaPinClient>,
        key_store: Arc<LocalKeyStore>,
    ) -> Self {
        let enforcer = Arc::new(DefaultToolInvocationEnforcer::new());
        Self {
            config,
            schema_pin,
            key_store,
            tools: Arc::new(RwLock::new(HashMap::new())),
            enforcer,
        }
    }

    /// Create a new secure MCP client with custom enforcer
    pub fn with_enforcer(
        config: McpClientConfig,
        schema_pin: Arc<dyn SchemaPinClient>,
        key_store: Arc<LocalKeyStore>,
        enforcer: Arc<dyn ToolInvocationEnforcer>,
    ) -> Self {
        Self {
            config,
            schema_pin,
            key_store,
            tools: Arc::new(RwLock::new(HashMap::new())),
            enforcer,
        }
    }

    /// Create a new secure MCP client with default components
    pub fn with_defaults(config: McpClientConfig) -> Result<Self, McpClientError> {
        let schema_pin = Arc::new(NativeSchemaPinClient::new());
        let key_store = Arc::new(LocalKeyStore::new()?);

        Ok(Self::new(config, schema_pin, key_store))
    }

    /// Verify the supplied schema with the exact trusted key, including a
    /// bounded first-use key fetch. Failed signature results are never promoted
    /// to Verified merely because the verification function returned Ok.
    async fn verify_schema(&self, tool: &McpTool) -> Result<VerificationStatus, McpClientError> {
        if self.config.verification_timeout_seconds == 0 {
            return Err(McpClientError::Timeout);
        }
        let schema =
            serde_json::to_vec(&tool.schema).map_err(|e| McpClientError::SerializationError {
                reason: e.to_string(),
            })?;
        let result = timeout(
            Duration::from_secs(self.config.verification_timeout_seconds),
            async {
                let key = self.fetch_and_pin_key(&tool.provider).await?;
                self.schema_pin
                    .verify_schema_bytes(&schema, &tool.name, &tool.provider.public_key_url, &key)
                    .await
                    .map_err(McpClientError::from)
            },
        )
        .await
        .map_err(|_| McpClientError::Timeout)?;
        Ok(match result {
            Ok(result) if result.success => VerificationStatus::Verified {
                result: Box::new(result),
                verified_at: chrono::Utc::now().to_rfc3339(),
            },
            Ok(result) => VerificationStatus::Failed {
                reason: result.message,
                failed_at: chrono::Utc::now().to_rfc3339(),
            },
            Err(error) => VerificationStatus::Failed {
                reason: error.to_string(),
                failed_at: chrono::Utc::now().to_rfc3339(),
            },
        })
    }

    /// Use the saved TOFU key as the verification anchor. A provider cannot
    /// replace that anchor by changing what its URL returns on a later fetch.
    /// Rotation requires an explicit update to the protected key store.
    async fn fetch_and_pin_key(&self, provider: &ToolProvider) -> Result<String, McpClientError> {
        if self.key_store.has_key(&provider.identifier)? {
            let key = self.key_store.get_key(&provider.identifier)?;
            if key.algorithm != "ES256" {
                return Err(McpClientError::KeyRetrievalFailed {
                    reason: "Pinned MCP key must use ES256".into(),
                });
            }
            return Ok(key.public_key);
        }
        let url = url::Url::parse(&provider.public_key_url).map_err(|e| {
            McpClientError::KeyFetchFailed {
                provider: provider.identifier.clone(),
                reason: e.to_string(),
            }
        })?;
        if url.scheme() != "https" {
            return Err(McpClientError::KeyFetchFailed {
                provider: provider.identifier.clone(),
                reason: "MCP key discovery requires HTTPS".into(),
            });
        }
        let key = NativeSchemaPinClient::new()
            .fetch_public_key(&provider.public_key_url)
            .await?;
        let fingerprint = hex::encode(sha2::Sha256::digest(key.as_bytes()));
        self.key_store.pin_key(PinnedKey::new(
            provider.identifier.clone(),
            key.clone(),
            "ES256".into(),
            fingerprint,
        ))?;
        Ok(key)
    }

    /// Check if a tool should be allowed based on verification status
    fn should_allow_tool(&self, tool: &McpTool) -> bool {
        match &tool.verification_status {
            VerificationStatus::Verified { .. } => tool.verification_status.is_verified(),
            VerificationStatus::Failed { .. } => false,
            VerificationStatus::Pending => !self.config.enforce_verification,
            VerificationStatus::Skipped { .. } => self.config.allow_unverified_in_dev,
        }
    }
}

#[async_trait]
impl McpClient for SecureMcpClient {
    async fn discover_tool(&self, mut tool: McpTool) -> Result<ToolDiscoveryEvent, McpClientError> {
        // Verify the tool's schema
        let verification_status = if self.config.enforce_verification {
            self.verify_schema(&tool).await?
        } else {
            VerificationStatus::Skipped {
                reason: "Verification disabled in configuration".to_string(),
            }
        };

        // Update tool with verification status
        tool.verification_status = verification_status;

        // Check if tool should be allowed (only if verification is enforced)
        if self.config.enforce_verification && !self.should_allow_tool(&tool) {
            return Err(McpClientError::VerificationFailed {
                reason: format!(
                    "Tool '{}' failed verification and cannot be added",
                    tool.name
                ),
            });
        }

        // Add tool to our collection
        {
            let mut tools = self.tools.write().await;
            tools.insert(tool.name.clone(), tool.clone());
        }

        Ok(ToolDiscoveryEvent {
            tool,
            source: "discovery".to_string(),
            discovered_at: chrono::Utc::now().to_rfc3339(),
        })
    }

    async fn get_tool(&self, name: &str) -> Result<McpTool, McpClientError> {
        let tools = self.tools.read().await;
        let tool = tools
            .get(name)
            .ok_or_else(|| McpClientError::ToolNotFound {
                name: name.to_string(),
            })?;

        // Check if verification is required
        if self.config.enforce_verification && !tool.verification_status.is_verified() {
            return Err(McpClientError::ToolNotVerified {
                name: name.to_string(),
            });
        }

        Ok(tool.clone())
    }

    async fn list_tools(&self) -> Result<Vec<McpTool>, McpClientError> {
        let tools = self.tools.read().await;
        Ok(tools.values().cloned().collect())
    }

    async fn list_verified_tools(&self) -> Result<Vec<McpTool>, McpClientError> {
        let tools = self.tools.read().await;
        Ok(tools
            .values()
            .filter(|tool| tool.verification_status.is_verified())
            .cloned()
            .collect())
    }

    async fn verify_tool(
        &self,
        request: ToolVerificationRequest,
    ) -> Result<ToolVerificationResponse, McpClientError> {
        let mut warnings = Vec::new();

        // Check if tool exists in our collection
        let tool_exists = {
            let tools = self.tools.read().await;
            tools.contains_key(&request.tool.name)
        };

        // If tool is already verified and not forcing re-verification, return current status
        if !request.force_reverify && tool_exists {
            let tools = self.tools.read().await;
            if let Some(existing_tool) = tools.get(&request.tool.name) {
                if existing_tool.verification_status.is_verified()
                    && existing_tool.schema == request.tool.schema
                    && existing_tool.provider.identifier == request.tool.provider.identifier
                    && existing_tool.provider.public_key_url == request.tool.provider.public_key_url
                    && existing_tool.provider.version == request.tool.provider.version
                {
                    warnings
                        .push("Tool already verified, use force_reverify to re-verify".to_string());
                    return Ok(ToolVerificationResponse {
                        tool_name: request.tool.name,
                        status: existing_tool.verification_status.clone(),
                        warnings,
                    });
                }
            }
        }

        // Perform verification
        let verification_status = self.verify_schema(&request.tool).await?;

        // Update tool in collection if it exists
        if tool_exists {
            let mut tools = self.tools.write().await;
            if let Some(existing_tool) = tools.get_mut(&request.tool.name) {
                // Store the exact contract that was verified, not a new status
                // attached to an older schema under the same tool name.
                *existing_tool = request.tool.clone();
                existing_tool.verification_status = verification_status.clone();
            }
        }

        Ok(ToolVerificationResponse {
            tool_name: request.tool.name,
            status: verification_status,
            warnings,
        })
    }

    async fn remove_tool(&self, name: &str) -> Result<Option<McpTool>, McpClientError> {
        let mut tools = self.tools.write().await;
        Ok(tools.remove(name))
    }

    async fn invoke_tool(
        &self,
        tool_name: &str,
        _arguments: serde_json::Value,
        context: InvocationContext,
    ) -> Result<InvocationResult, McpClientError> {
        // Get the tool first
        let tool = self.get_tool(tool_name).await?;

        // Execute with enforcement
        self.enforcer
            .execute_tool_with_enforcement(&tool, context)
            .await
            .map_err(|e| match e {
                ToolInvocationError::InvocationBlocked {
                    tool_name,
                    reason: _,
                } => McpClientError::ToolNotVerified { name: tool_name },
                ToolInvocationError::ToolNotFound { tool_name } => {
                    McpClientError::ToolNotFound { name: tool_name }
                }
                ToolInvocationError::VerificationRequired { tool_name, .. } => {
                    McpClientError::ToolNotVerified { name: tool_name }
                }
                ToolInvocationError::VerificationFailed {
                    tool_name: _,
                    reason,
                } => McpClientError::VerificationFailed { reason },
                _ => McpClientError::CommunicationError {
                    reason: e.to_string(),
                },
            })
    }
}

/// Mock MCP client for testing
pub struct MockMcpClient {
    tools: Arc<RwLock<HashMap<String, McpTool>>>,
    should_verify_successfully: bool,
}

impl MockMcpClient {
    /// Create a new mock client that succeeds verification
    pub fn new_success() -> Self {
        Self {
            tools: Arc::new(RwLock::new(HashMap::new())),
            should_verify_successfully: true,
        }
    }

    /// Create a new mock client that fails verification
    pub fn new_failure() -> Self {
        Self {
            tools: Arc::new(RwLock::new(HashMap::new())),
            should_verify_successfully: false,
        }
    }
}

#[async_trait]
impl McpClient for MockMcpClient {
    async fn discover_tool(&self, mut tool: McpTool) -> Result<ToolDiscoveryEvent, McpClientError> {
        // Mock verification
        tool.verification_status = if self.should_verify_successfully {
            VerificationStatus::Verified {
                result: Box::new(crate::integrations::schemapin::VerificationResult {
                    success: true,
                    message: "Mock verification successful".to_string(),
                    schema_hash: Some("mock_hash".to_string()),
                    public_key_url: Some(tool.provider.public_key_url.clone()),
                    signature: None,
                    metadata: None,
                    timestamp: Some(chrono::Utc::now().to_rfc3339()),
                }),
                verified_at: chrono::Utc::now().to_rfc3339(),
            }
        } else {
            VerificationStatus::Failed {
                reason: "Mock verification failed".to_string(),
                failed_at: chrono::Utc::now().to_rfc3339(),
            }
        };

        if !self.should_verify_successfully {
            return Err(McpClientError::VerificationFailed {
                reason: "Mock verification failed".to_string(),
            });
        }

        let mut tools = self.tools.write().await;
        tools.insert(tool.name.clone(), tool.clone());

        Ok(ToolDiscoveryEvent {
            tool,
            source: "mock".to_string(),
            discovered_at: chrono::Utc::now().to_rfc3339(),
        })
    }

    async fn get_tool(&self, name: &str) -> Result<McpTool, McpClientError> {
        let tools = self.tools.read().await;
        tools
            .get(name)
            .cloned()
            .ok_or_else(|| McpClientError::ToolNotFound {
                name: name.to_string(),
            })
    }

    async fn list_tools(&self) -> Result<Vec<McpTool>, McpClientError> {
        let tools = self.tools.read().await;
        Ok(tools.values().cloned().collect())
    }

    async fn list_verified_tools(&self) -> Result<Vec<McpTool>, McpClientError> {
        let tools = self.tools.read().await;
        Ok(tools
            .values()
            .filter(|tool| tool.verification_status.is_verified())
            .cloned()
            .collect())
    }

    async fn verify_tool(
        &self,
        request: ToolVerificationRequest,
    ) -> Result<ToolVerificationResponse, McpClientError> {
        let status = if self.should_verify_successfully {
            VerificationStatus::Verified {
                result: Box::new(crate::integrations::schemapin::VerificationResult {
                    success: true,
                    message: "Mock verification successful".to_string(),
                    schema_hash: Some("mock_hash".to_string()),
                    public_key_url: Some(request.tool.provider.public_key_url.clone()),
                    signature: None,
                    metadata: None,
                    timestamp: Some(chrono::Utc::now().to_rfc3339()),
                }),
                verified_at: chrono::Utc::now().to_rfc3339(),
            }
        } else {
            VerificationStatus::Failed {
                reason: "Mock verification failed".to_string(),
                failed_at: chrono::Utc::now().to_rfc3339(),
            }
        };

        Ok(ToolVerificationResponse {
            tool_name: request.tool.name,
            status,
            warnings: vec![],
        })
    }

    async fn remove_tool(&self, name: &str) -> Result<Option<McpTool>, McpClientError> {
        let mut tools = self.tools.write().await;
        Ok(tools.remove(name))
    }

    async fn invoke_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
        _context: InvocationContext,
    ) -> Result<InvocationResult, McpClientError> {
        // Get the tool to check its verification status
        let tool = self.get_tool(tool_name).await?;

        // Mock enforcement - only allow verified tools if should_verify_successfully is true
        if !self.should_verify_successfully && !tool.verification_status.is_verified() {
            return Err(McpClientError::ToolNotVerified {
                name: tool_name.to_string(),
            });
        }

        // Return mock successful result
        Ok(InvocationResult {
            success: true,
            result: serde_json::json!({
                "status": "success",
                "tool": tool_name,
                "arguments": arguments
            }),
            execution_time: Duration::from_millis(50),
            warnings: vec![],
            metadata: std::collections::HashMap::new(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::integrations::schemapin::types::KeyStoreConfig;
    use crate::integrations::schemapin::MockNativeSchemaPinClient;

    fn create_test_tool() -> McpTool {
        McpTool {
            name: "test_tool".to_string(),
            description: "A test tool".to_string(),
            schema: serde_json::json!({
                "type": "object",
                "properties": {
                    "input": {"type": "string"}
                }
            }),
            provider: ToolProvider {
                identifier: "example.com".to_string(),
                name: "Example Provider".to_string(),
                public_key_url: "https://example.com/pubkey".to_string(),
                version: Some("1.0.0".to_string()),
            },
            verification_status: VerificationStatus::Pending,
            metadata: None,
            sensitive_params: vec![],
        }
    }

    /// Create an isolated key store backed by a temp directory so tests
    /// never collide with each other or with the user's real key store.
    fn create_temp_key_store() -> (LocalKeyStore, tempfile::TempDir) {
        let temp_dir = tempfile::TempDir::new().unwrap();
        let store_path = temp_dir.path().join("test_keys.json");
        let config = KeyStoreConfig {
            store_path,
            create_if_missing: true,
            file_permissions: Some(0o600),
            closed_world: false,
        };
        let store = LocalKeyStore::with_config(config).unwrap();
        (store, temp_dir)
    }

    fn sign_test_tool(tool: &mut McpTool, private_key: &str) {
        tool.schema.as_object_mut().unwrap().remove("signature");
        let bytes = schemapin::canonicalize::canonicalize_schema(&tool.schema).into_bytes();
        tool.schema["signature"] =
            serde_json::json!(schemapin::crypto::sign_data(private_key, &bytes).unwrap());
    }

    #[tokio::test]
    async fn native_verification_requires_a_signature_from_the_pinned_key() {
        let key = schemapin::crypto::generate_key_pair().unwrap();
        let substituted = schemapin::crypto::generate_key_pair().unwrap();
        let (store, _directory) = create_temp_key_store();
        store
            .pin_key(PinnedKey::new(
                "example.com".into(),
                key.public_key_pem.clone(),
                "ES256".into(),
                hex::encode(sha2::Sha256::digest(key.public_key_pem.as_bytes())),
            ))
            .unwrap();
        let client = SecureMcpClient::new(
            McpClientConfig::default(),
            Arc::new(NativeSchemaPinClient::new()),
            Arc::new(store),
        );
        let mut tool = create_test_tool();
        assert!(client.discover_tool(tool.clone()).await.is_err());
        sign_test_tool(&mut tool, &substituted.private_key_pem);
        assert!(client.discover_tool(tool.clone()).await.is_err());
        sign_test_tool(&mut tool, &key.private_key_pem);
        assert!(client
            .discover_tool(tool.clone())
            .await
            .unwrap()
            .tool
            .verification_status
            .is_verified());
        tool.schema["properties"]["input"]["type"] = serde_json::json!("integer");
        assert!(client.discover_tool(tool).await.is_err());
    }

    #[tokio::test]
    async fn verification_cache_and_registry_bind_status_to_the_exact_schema() {
        let key = schemapin::crypto::generate_key_pair().unwrap();
        let (store, _directory) = create_temp_key_store();
        store
            .pin_key(PinnedKey::new(
                "example.com".into(),
                key.public_key_pem.clone(),
                "ES256".into(),
                hex::encode(sha2::Sha256::digest(key.public_key_pem.as_bytes())),
            ))
            .unwrap();
        let client = SecureMcpClient::new(
            McpClientConfig::default(),
            Arc::new(NativeSchemaPinClient::new()),
            Arc::new(store),
        );
        let mut tool = create_test_tool();
        sign_test_tool(&mut tool, &key.private_key_pem);
        client.discover_tool(tool.clone()).await.unwrap();
        tool.schema["properties"]["input"]["type"] = serde_json::json!("integer");
        let response = client
            .verify_tool(ToolVerificationRequest {
                tool: tool.clone(),
                force_reverify: false,
            })
            .await
            .unwrap();
        assert!(response.status.is_failed());
        assert!(client.get_tool(&tool.name).await.is_err());
        sign_test_tool(&mut tool, &key.private_key_pem);
        let response = client
            .verify_tool(ToolVerificationRequest {
                tool: tool.clone(),
                force_reverify: false,
            })
            .await
            .unwrap();
        assert!(response.status.is_verified());
        assert_eq!(
            client.get_tool(&tool.name).await.unwrap().schema,
            tool.schema
        );
    }

    #[tokio::test]
    async fn test_mock_client_success() {
        let client = MockMcpClient::new_success();
        let tool = create_test_tool();

        let event = client.discover_tool(tool.clone()).await.unwrap();
        assert!(event.tool.verification_status.is_verified());

        let retrieved_tool = client.get_tool(&tool.name).await.unwrap();
        assert_eq!(retrieved_tool.name, tool.name);
    }

    #[tokio::test]
    async fn test_mock_client_failure() {
        let client = MockMcpClient::new_failure();
        let tool = create_test_tool();

        let result = client.discover_tool(tool).await;
        assert!(result.is_err());
        assert!(matches!(
            result,
            Err(McpClientError::VerificationFailed { .. })
        ));
    }

    #[tokio::test]
    async fn test_secure_client_with_mock_components() {
        let config = McpClientConfig::default();
        let schema_pin = Arc::new(MockNativeSchemaPinClient::new_success());
        let (key_store, _temp_dir) = create_temp_key_store();

        // Pre-pin so the real HTTPS fetch is skipped in tests
        key_store
            .pin_key(PinnedKey::new(
                "example.com".to_string(),
                "test_key".to_string(),
                "ES256".to_string(),
                "test_fingerprint".to_string(),
            ))
            .unwrap();

        let key_store = Arc::new(key_store);

        let client = SecureMcpClient::new(config, schema_pin, key_store);
        let tool = create_test_tool();

        let event = client.discover_tool(tool.clone()).await.unwrap();
        assert!(event.tool.verification_status.is_verified());

        let tools = client.list_verified_tools().await.unwrap();
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0].name, tool.name);
    }

    #[tokio::test]
    async fn test_verification_enforcement() {
        let config = McpClientConfig {
            enforce_verification: true,
            ..Default::default()
        };

        let schema_pin = Arc::new(MockNativeSchemaPinClient::new_failure());
        let (key_store, _temp_dir) = create_temp_key_store();

        // Pre-pin so the real HTTPS fetch is skipped in tests
        key_store
            .pin_key(PinnedKey::new(
                "example.com".to_string(),
                "test_key".to_string(),
                "ES256".to_string(),
                "test_fingerprint".to_string(),
            ))
            .unwrap();

        let key_store = Arc::new(key_store);

        let client = SecureMcpClient::new(config, schema_pin, key_store);
        let tool = create_test_tool();

        let result = client.discover_tool(tool).await;
        assert!(result.is_err());
        assert!(matches!(
            result,
            Err(McpClientError::VerificationFailed { .. })
        ));
    }
}