MockLoop MCP Server

#!/usr/bin/env python3 """ Advanced SchemaPin Usage Example This example demonstrates advanced SchemaPin integration scenarios including: - Custom verification workflows - Integration with MCP proxy - Batch operations - Error handling and recovery - Performance monitoring """ import asyncio import json import sys import time from pathlib import Path from typing import Any from unittest.mock import patch # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src")) from mockloop_mcp.schemapin import ( SchemaPinConfig, SchemaVerificationInterceptor, KeyPinningManager, PolicyHandler, SchemaPinAuditLogger, VerificationResult, PolicyAction, PolicyDecision ) class SchemaPinWorkflowManager: """Advanced workflow manager for SchemaPin operations.""" def __init__(self, config: SchemaPinConfig): self.config = config self.interceptor = SchemaVerificationInterceptor(config) self.key_manager = KeyPinningManager(config.key_pin_storage_path) self.policy_handler = PolicyHandler(config) self.audit_logger = SchemaPinAuditLogger() self.performance_metrics = {} async def verify_tool_batch(self, tools: list[dict[str, Any]]) -> list[VerificationResult]: """Verify multiple tools in batch with performance tracking.""" start_time = time.time() results = [] print(f"Starting batch verification of {len(tools)} tools...") # Process tools concurrently for better performance tasks = [] for tool in tools: task = self.interceptor.verify_tool_schema( tool_name=tool["name"], schema=tool["schema"], signature=tool.get("signature"), domain=tool.get("domain") ) tasks.append(task) results = await asyncio.gather(*tasks, return_exceptions=True) # Process results and handle exceptions processed_results = [] for i, result in enumerate(results): if isinstance(result, Exception): error_result = VerificationResult( valid=False, tool_id=tools[i]["name"], error=str(result) ) processed_results.append(error_result) else: processed_results.append(result) execution_time = (time.time() - start_time) * 1000 self.performance_metrics["batch_verification"] = { "tools_count": len(tools), "execution_time_ms": execution_time, "avg_time_per_tool": execution_time / len(tools) } print(f"Batch verification completed in {execution_time:.2f}ms") return processed_results async def setup_trusted_domain(self, domain: str, public_key: str, tools: list[str]) -> dict[str, bool]: """Set up a trusted domain with multiple tools.""" print(f"Setting up trusted domain: {domain}") results = {} for tool_name in tools: tool_id = f"{domain}/{tool_name}" success = self.key_manager.pin_key( tool_id=tool_id, domain=domain, public_key_pem=public_key, metadata={ "setup_type": "trusted_domain", "setup_time": time.time(), "tool_name": tool_name } ) results[tool_name] = success if success: await self.audit_logger.log_key_pinning_event( tool_id, domain, public_key, "pin" ) print(f"Domain setup complete: {sum(results.values())}/{len(tools)} tools pinned") return results async def handle_verification_failure(self, result: VerificationResult, tool_name: str) -> PolicyDecision: """Advanced failure handling with custom recovery strategies.""" print(f"Handling verification failure for {tool_name}") # Get policy decision decision = await self.policy_handler.evaluate_verification_result(result, tool_name) # Log the decision await self.audit_logger.log_policy_decision( result.tool_id, decision.action.value, decision.reason, decision.policy_mode ) # Implement custom recovery strategies if decision.action == PolicyAction.BLOCK: print(f"🚫 BLOCKED: {tool_name} - {decision.reason}") # Check if we can attempt key recovery if result.domain and "key" in result.error.lower(): print("Attempting key recovery...") recovered_key = await self._attempt_key_recovery(result.domain) if recovered_key: print("✓ Key recovery successful, retrying verification") # In a real implementation, you'd retry verification here elif decision.action == PolicyAction.WARN: print(f"⚠️ WARNING: {tool_name} - {decision.reason}") print("Tool execution will proceed with monitoring") elif decision.action == PolicyAction.LOG: print(f"📝 LOGGED: {tool_name} - {decision.reason}") return decision async def _attempt_key_recovery(self, domain: str) -> str | None: """Attempt to recover a public key through alternative methods.""" # Try alternative discovery endpoints alternative_endpoints = [ f"https://{domain}/security/schemapin.json", f"https://{domain}/api/v1/schemapin", f"https://security.{domain}/schemapin.json" ] for endpoint in alternative_endpoints: try: # In a real implementation, you'd make HTTP requests here print(f"Trying alternative endpoint: {endpoint}") # Simulate discovery attempt await asyncio.sleep(0.1) except Exception as e: # Log the exception for debugging print(f"Failed to connect to {endpoint}: {e}") continue return None def get_performance_report(self) -> dict[str, Any]: """Generate performance report.""" return { "metrics": self.performance_metrics, "key_stats": { "total_pinned_keys": len(self.key_manager.list_pinned_keys()), "audit_stats": self.audit_logger.get_verification_stats() } } async def advanced_verification_workflow(): """Demonstrate advanced verification workflows.""" print("=== Advanced Verification Workflow ===\n") # Configure for production-like scenario config = SchemaPinConfig( policy_mode="enforce", auto_pin_keys=False, key_pin_storage_path="advanced_keys.db", trusted_domains=["api.corp.com", "tools.internal.com"], discovery_timeout=30, interactive_mode=False ) workflow_manager = SchemaPinWorkflowManager(config) # Set up trusted domain with multiple tools trusted_domain = "api.corp.com" trusted_key = """-----BEGIN PUBLIC KEY----- MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAETrustedKey1234567890abcdefgh ijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890trusted -----END PUBLIC KEY-----""" trusted_tools = ["database_query", "file_operations", "api_gateway", "auth_service"] setup_results = await workflow_manager.setup_trusted_domain( trusted_domain, trusted_key, trusted_tools ) print(f"✓ Trusted domain setup: {setup_results}\n") # Prepare batch of tools for verification tools_batch = [ { "name": "database_query", "domain": "api.corp.com", "schema": { "name": "database_query", "description": "Execute database queries", "parameters": {"type": "object"} }, "signature": "valid_signature_1" }, { "name": "file_operations", "domain": "api.corp.com", "schema": { "name": "file_operations", "description": "File system operations", "parameters": {"type": "object"} }, "signature": "valid_signature_2" }, { "name": "suspicious_tool", "domain": "malicious.com", "schema": { "name": "suspicious_tool", "description": "Suspicious operations", "parameters": {"type": "object"} }, "signature": "invalid_signature" }, { "name": "unsigned_tool", "domain": "unknown.com", "schema": { "name": "unsigned_tool", "description": "Tool without signature", "parameters": {"type": "object"} }, "signature": None } ] # Mock signature verification for demonstration def mock_verify_signature(schema, signature, public_key): # Simulate verification logic if signature and "valid_signature" in signature: return True return False with patch.object(workflow_manager.interceptor, '_verify_signature', side_effect=mock_verify_signature): # Perform batch verification results = await workflow_manager.verify_tool_batch(tools_batch) print("--- Batch Verification Results ---") for i, result in enumerate(results): tool = tools_batch[i] print(f"Tool: {tool['name']}") print(f" Valid: {result.valid}") print(f" Domain: {result.domain}") print(f" Key pinned: {result.key_pinned}") if result.error: print(f" Error: {result.error}") # Handle failures with advanced strategies if not result.valid: await workflow_manager.handle_verification_failure(result, tool['name']) print() # Generate performance report report = workflow_manager.get_performance_report() print("--- Performance Report ---") print(json.dumps(report, indent=2, default=str)) async def mcp_proxy_integration_example(): """Demonstrate integration with MCP proxy patterns.""" print("\n=== MCP Proxy Integration Example ===\n") class SchemaPinMCPProxy: """MCP Proxy with SchemaPin integration.""" def __init__(self, config: SchemaPinConfig): self.config = config self.interceptor = SchemaVerificationInterceptor(config) self.request_cache = {} async def proxy_tool_request(self, tool_name: str, tool_schema: dict[str, Any], signature: str | None, domain: str | None, request_data: dict[str, Any]) -> dict[str, Any]: """Proxy tool request with SchemaPin verification.""" print(f"Proxying request for tool: {tool_name}") # Verify schema first verification_result = await self.interceptor.verify_tool_schema( tool_name, tool_schema, signature, domain ) if not verification_result.valid: return { "error": "Schema verification failed", "details": verification_result.error, "tool_id": verification_result.tool_id } # Cache verification result for performance cache_key = f"{domain}/{tool_name}" self.request_cache[cache_key] = { "verified_at": time.time(), "result": verification_result } # Simulate tool execution print(f"✓ Schema verified, executing tool: {tool_name}") return { "success": True, "tool_id": verification_result.tool_id, "verified": True, "key_pinned": verification_result.key_pinned, "result": f"Executed {tool_name} with data: {request_data}" } def get_verification_cache_stats(self) -> dict[str, Any]: """Get cache statistics.""" return { "cached_verifications": len(self.request_cache), "cache_entries": list(self.request_cache.keys()) } # Initialize proxy proxy_config = SchemaPinConfig( policy_mode="warn", cache_ttl=3600, key_pin_storage_path="proxy_keys.db" ) proxy = SchemaPinMCPProxy(proxy_config) # Simulate proxy requests requests = [ { "tool_name": "secure_api_call", "domain": "api.secure.com", "schema": {"name": "secure_api_call", "description": "Secure API operations"}, "signature": "proxy_signature_1", "request_data": {"endpoint": "/api/v1/data", "method": "GET"} }, { "tool_name": "data_processor", "domain": "processing.com", "schema": {"name": "data_processor", "description": "Process data"}, "signature": None, # Unsigned tool "request_data": {"data": [1, 2, 3, 4, 5]} } ] for request in requests: result = await proxy.proxy_tool_request( request["tool_name"], request["schema"], request["signature"], request["domain"], request["request_data"] ) print(f"Proxy result for {request['tool_name']}:") print(json.dumps(result, indent=2)) print() # Show cache stats cache_stats = proxy.get_verification_cache_stats() print("Proxy cache statistics:") print(json.dumps(cache_stats, indent=2)) async def error_recovery_and_monitoring(): """Demonstrate error recovery and monitoring capabilities.""" print("\n=== Error Recovery and Monitoring ===\n") class SchemaPinMonitor: """Monitor SchemaPin operations and handle errors.""" def __init__(self, config: SchemaPinConfig): self.config = config self.interceptor = SchemaVerificationInterceptor(config) self.audit_logger = SchemaPinAuditLogger() self.error_counts = {} self.recovery_attempts = {} async def monitored_verification(self, tool_name: str, schema: dict[str, Any], signature: str | None, domain: str | None) -> dict[str, Any]: """Perform verification with monitoring and recovery.""" tool_id = f"{domain}/{tool_name}" if domain else tool_name try: result = await self.interceptor.verify_tool_schema( tool_name, schema, signature, domain ) if not result.valid: await self._handle_verification_error(tool_id, result.error) return { "verification_result": result, "monitoring_data": { "error_count": self.error_counts.get(tool_id, 0), "recovery_attempts": self.recovery_attempts.get(tool_id, 0) } } except Exception as e: await self._handle_system_error(tool_id, str(e)) return { "verification_result": VerificationResult( valid=False, tool_id=tool_id, error=f"System error: {e}" ), "monitoring_data": { "system_error": True, "error_count": self.error_counts.get(tool_id, 0) } } async def _handle_verification_error(self, tool_id: str, error: str): """Handle verification errors with recovery strategies.""" self.error_counts[tool_id] = self.error_counts.get(tool_id, 0) + 1 print(f"⚠️ Verification error for {tool_id}: {error}") # Implement recovery strategies based on error type if "key" in error.lower() and "not found" in error.lower(): await self._attempt_key_rediscovery(tool_id) elif "signature" in error.lower(): await self._log_signature_failure(tool_id, error) elif "timeout" in error.lower(): await self._handle_timeout_error(tool_id) async def _handle_system_error(self, tool_id: str, error: str): """Handle system-level errors.""" print(f"🚨 System error for {tool_id}: {error}") await self.audit_logger.log_verification_error(tool_id, None, f"System: {error}") async def _attempt_key_rediscovery(self, tool_id: str): """Attempt to rediscover keys for failed tools.""" self.recovery_attempts[tool_id] = self.recovery_attempts.get(tool_id, 0) + 1 print(f"🔄 Attempting key rediscovery for {tool_id} (attempt #{self.recovery_attempts[tool_id]})") # In a real implementation, this would attempt actual key rediscovery await asyncio.sleep(0.1) # Simulate recovery attempt async def _log_signature_failure(self, tool_id: str, error: str): """Log signature verification failures for security analysis.""" await self.audit_logger.log_verification_error(tool_id, None, f"Signature failure: {error}") print(f"🔐 Signature failure logged for security analysis: {tool_id}") async def _handle_timeout_error(self, tool_id: str): """Handle timeout errors with retry logic.""" print(f"⏱️ Timeout error for {tool_id}, implementing backoff strategy") # In a real implementation, this would implement exponential backoff def get_monitoring_summary(self) -> dict[str, Any]: """Get monitoring summary.""" return { "total_errors": sum(self.error_counts.values()), "tools_with_errors": len(self.error_counts), "total_recovery_attempts": sum(self.recovery_attempts.values()), "error_breakdown": self.error_counts.copy(), "recovery_breakdown": self.recovery_attempts.copy() } # Initialize monitor monitor_config = SchemaPinConfig( policy_mode="log", discovery_timeout=5, # Short timeout to trigger errors key_pin_storage_path="monitor_keys.db" ) monitor = SchemaPinMonitor(monitor_config) # Test various error scenarios test_scenarios = [ { "name": "valid_tool", "domain": "valid.com", "schema": {"name": "valid_tool"}, "signature": "valid_sig", "expected": "success" }, { "name": "missing_key_tool", "domain": "unknown.com", "schema": {"name": "missing_key_tool"}, "signature": "some_sig", "expected": "key_not_found" }, { "name": "invalid_signature_tool", "domain": "malicious.com", "schema": {"name": "invalid_signature_tool"}, "signature": "bad_signature", "expected": "signature_failure" }, { "name": "unsigned_tool", "domain": "unsigned.com", "schema": {"name": "unsigned_tool"}, "signature": None, "expected": "no_signature" } ] print("Testing error scenarios...") for scenario in test_scenarios: print(f"\n--- Testing: {scenario['name']} ---") result = await monitor.monitored_verification( scenario["name"], scenario["schema"], scenario["signature"], scenario["domain"] ) print(f"Valid: {result['verification_result'].valid}") if result['verification_result'].error: print(f"Error: {result['verification_result'].error}") print(f"Monitoring: {result['monitoring_data']}") # Show monitoring summary summary = monitor.get_monitoring_summary() print("\n--- Monitoring Summary ---") print(json.dumps(summary, indent=2)) async def main(): """Run all advanced examples.""" print("Advanced SchemaPin Integration Examples") print("=" * 60) try: await advanced_verification_workflow() await mcp_proxy_integration_example() await error_recovery_and_monitoring() print("\n" + "=" * 60) print("✓ All advanced examples completed successfully!") print("\nKey takeaways:") print("1. Batch verification improves performance for multiple tools") print("2. MCP proxy integration provides seamless security") print("3. Error recovery and monitoring ensure robust operations") print("4. Performance metrics help optimize SchemaPin usage") except Exception as e: print(f"\n❌ Error running advanced examples: {e}") import traceback traceback.print_exc() finally: # Clean up example files cleanup_files = [ "advanced_keys.db", "proxy_keys.db", "monitor_keys.db" ] print("\nCleaning up example files...") for file_path in cleanup_files: try: Path(file_path).unlink(missing_ok=True) print(f"✓ Removed {file_path}") except Exception as e: print(f"⚠ Could not remove {file_path}: {e}") if __name__ == "__main__": asyncio.run(main())