mcp-proxy-adapter 2.0.1__py3-none-any.whl → 6.9.50__py3-none-any.whl

mcp_proxy_adapter/examples/universal_client.py

@@ -0,0 +1,674 @@
+"""
+Universal Client Example
+This module demonstrates all possible secure connection methods to the server
+using the mcp_security_framework. The client supports all authentication methods
+and connection types supported by the security framework.
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import asyncio
+import json
+import os
+import ssl
+import time
+from pathlib import Path
+import aiohttp
+import requests
+
+# Import security framework components
+try:
+    from mcp_security_framework import (
+        SecurityManager,
+        AuthManager,
+        CertificateManager,
+        PermissionManager,
+    )
+        generate_api_key,
+        create_jwt_token,
+        validate_jwt_token,
+    )
+        extract_roles_from_cert,
+        validate_certificate_chain,
+    )
+        create_ssl_context,
+        validate_server_certificate,
+    )
+
+    SECURITY_FRAMEWORK_AVAILABLE = True
+except ImportError:
+    SECURITY_FRAMEWORK_AVAILABLE = False
+    print("Warning: mcp_security_framework not available. Using basic HTTP client.")
+
+
+class UniversalClient:
+    """
+    Universal client that demonstrates all possible secure connection methods.
+    Supports:
+    - HTTP/HTTPS connections
+    - API Key authentication
+    - JWT token authentication
+    - Certificate-based authentication
+    - SSL/TLS with custom certificates
+    - Role-based access control
+    - Rate limiting awareness
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initialize universal client with configuration.
+        Args:
+            config: Client configuration with security settings
+        """
+        self.config = config
+        self.base_url = config.get("server_url", "http://localhost:8000")
+        self.timeout = config.get("timeout", 30)
+        self.retry_attempts = config.get("retry_attempts", 3)
+        self.retry_delay = config.get("retry_delay", 1)
+        # Security configuration
+        self.security_config = config.get("security", {})
+        self.auth_method = self.security_config.get("auth_method", "none")
+        # Initialize security managers if framework is available
+        self.security_manager = None
+        self.auth_manager = None
+        self.cert_manager = None
+        if SECURITY_FRAMEWORK_AVAILABLE:
+            self._initialize_security_managers()
+        # Session management
+        self.session: Optional[aiohttp.ClientSession] = None
+        self.current_token: Optional[str] = None
+        self.token_expiry: Optional[float] = None
+        print(f"Universal client initialized with auth method: {self.auth_method}")
+
+    def _initialize_security_managers(self) -> None:
+        """Initialize security framework managers."""
+        try:
+            # Initialize security manager
+            self.security_manager = SecurityManager(self.security_config)
+            # Initialize permission manager first
+            permissions_config = self.security_config.get("permissions", {})
+            self.permission_manager = PermissionManager(permissions_config)
+            # Initialize auth manager with permission_manager
+            auth_config = self.security_config.get("auth", {})
+            self.auth_manager = AuthManager(auth_config, self.permission_manager)
+            # Initialize certificate manager
+            cert_config = self.security_config.get("certificates", {})
+            self.cert_manager = CertificateManager(cert_config)
+            print("Security framework managers initialized successfully")
+        except Exception as e:
+            print(f"Warning: Failed to initialize security managers: {e}")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.disconnect()
+
+    async def connect(self) -> None:
+        """Establish connection with authentication."""
+        print(
+            f"Connecting to {self.base_url} with {self.auth_method} authentication..."
+        )
+        # Create SSL context
+        ssl_context = self._create_ssl_context()
+        # Create connector with SSL context
+        connector = None
+        if ssl_context:
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+        # Create session
+        self.session = aiohttp.ClientSession(connector=connector)
+        # Perform authentication based on method
+        if self.auth_method == "api_key":
+            await self._authenticate_api_key()
+        elif self.auth_method == "jwt":
+            await self._authenticate_jwt()
+        elif self.auth_method == "certificate":
+            await self._authenticate_certificate()
+        elif self.auth_method == "basic":
+            await self._authenticate_basic()
+        else:
+            print("No authentication required")
+        print("Connection established successfully")
+
+    async def disconnect(self) -> None:
+        """Close connection and cleanup."""
+        if self.session:
+            await self.session.close()
+            self.session = None
+        print("Connection closed")
+
+    async def _authenticate_api_key(self) -> None:
+        """Authenticate using API key."""
+        api_key_config = self.security_config.get("api_key", {})
+        api_key = api_key_config.get("key")
+        if not api_key:
+            raise ValueError("API key not provided in configuration")
+        # Store API key for requests
+        self.current_token = api_key
+        print(f"Authenticated with API key: {api_key[:8]}...")
+
+    async def _authenticate_jwt(self) -> None:
+        """Authenticate using JWT token."""
+        jwt_config = self.security_config.get("jwt", {})
+        # Check if we have a stored token that's still valid
+        if self.current_token and self.token_expiry and time.time() < self.token_expiry:
+            print("Using existing JWT token")
+            return
+        # Get credentials for JWT
+        username = jwt_config.get("username")
+        password = jwt_config.get("password")
+        secret = jwt_config.get("secret")
+        if not all([username, password, secret]):
+            raise ValueError("JWT credentials not provided in configuration")
+        # Create JWT token
+        if SECURITY_FRAMEWORK_AVAILABLE:
+            self.current_token = create_jwt_token(
+                username, secret, expiry_hours=jwt_config.get("expiry_hours", 24)
+            )
+        else:
+            # Simple JWT creation (for demonstration)
+            import jwt
+
+            payload = {
+                "username": username,
+                "exp": time.time() + (jwt_config.get("expiry_hours", 24) * 3600),
+            }
+            self.current_token = jwt.encode(payload, secret, algorithm="HS256")
+        self.token_expiry = time.time() + (jwt_config.get("expiry_hours", 24) * 3600)
+        print(f"Authenticated with JWT token: {self.current_token[:20]}...")
+
+    async def _authenticate_certificate(self) -> None:
+        """Authenticate using client certificate."""
+        cert_config = self.security_config.get("certificate", {})
+        cert_file = cert_config.get("cert_file")
+        key_file = cert_config.get("key_file")
+        if not cert_file or not key_file:
+            raise ValueError("Certificate files not provided in configuration")
+        # Validate certificate
+        if SECURITY_FRAMEWORK_AVAILABLE and self.cert_manager:
+            try:
+                cert_info = self.cert_manager.validate_certificate(cert_file, key_file)
+                print(f"Certificate validated: {cert_info.get('subject', 'Unknown')}")
+                # Extract roles from certificate
+                roles = extract_roles_from_cert(cert_file)
+                if roles:
+                    print(f"Certificate roles: {roles}")
+            except Exception as e:
+                print(f"Warning: Certificate validation failed: {e}")
+        print("Certificate authentication prepared")
+
+    async def _authenticate_basic(self) -> None:
+        """Authenticate using basic authentication."""
+        basic_config = self.security_config.get("basic", {})
+        username = basic_config.get("username")
+        password = basic_config.get("password")
+        if not username or not password:
+            raise ValueError("Basic auth credentials not provided in configuration")
+        import base64
+
+        credentials = base64.b64encode(f"{username}:{password}".encode()).decode()
+        self.current_token = f"Basic {credentials}"
+        print(f"Authenticated with basic auth: {username}")
+
+    def _get_auth_headers(self) -> Dict[str, str]:
+        """Get authentication headers for requests."""
+        headers = {"Content-Type": "application/json"}
+        if not self.current_token:
+            return headers
+        if self.auth_method == "api_key":
+            api_key_config = self.security_config.get("api_key", {})
+            header_name = api_key_config.get("header", "X-API-Key")
+            headers[header_name] = self.current_token
+        elif self.auth_method == "jwt":
+            headers["Authorization"] = f"Bearer {self.current_token}"
+        elif self.auth_method == "basic":
+            headers["Authorization"] = self.current_token
+        return headers
+
+    def _create_ssl_context(self) -> Optional[ssl.SSLContext]:
+        """Create SSL context for secure connections."""
+        ssl_config = self.security_config.get("ssl", {})
+        if not ssl_config.get("enabled", False):
+            return None
+        try:
+            ssl_config = self.security_config.get("ssl", {})
+            if not ssl_config.get("enabled", False):
+                return None
+            # Create SSL context using security framework
+            if self.security_manager:
+                return self.security_manager.create_client_ssl_context()
+            # Fallback SSL context creation
+            context = ssl.create_default_context()
+            # Add client certificate if provided
+            cert_config = self.security_config.get("certificate", {})
+            if cert_config.get("enabled", False):
+                cert_file = cert_config.get("cert_file")
+                key_file = cert_config.get("key_file")
+                if cert_file and key_file:
+                    context.load_cert_chain(cert_file, key_file)
+            # Add CA certificate if provided
+            ca_cert_file = ssl_config.get("ca_cert_file") or ssl_config.get("ca_cert")
+            if ca_cert_file and os.path.exists(ca_cert_file):
+                context.load_verify_locations(ca_cert_file)
+            # Configure verification
+            if ssl_config.get("check_hostname", True):
+                context.check_hostname = True
+                context.verify_mode = ssl.CERT_REQUIRED
+            else:
+                context.check_hostname = False
+                context.verify_mode = ssl.CERT_NONE
+            return context
+        except Exception as e:
+            print(f"Warning: Failed to create SSL context: {e}")
+            return None
+
+    async def request(
+        self,
+        method: str,
+        endpoint: str,
+        data: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Make authenticated request to server.
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            endpoint: API endpoint
+            data: Request data
+            headers: Additional headers
+        Returns:
+            Response data
+        """
+        url = urljoin(self.base_url, endpoint)
+        # Prepare headers
+        request_headers = self._get_auth_headers()
+        if headers:
+            request_headers.update(headers)
+        try:
+            for attempt in range(self.retry_attempts):
+                try:
+                    async with self.session.request(
+                        method,
+                        url,
+                        json=data,
+                        headers=request_headers,
+                        timeout=aiohttp.ClientTimeout(total=self.timeout),
+                    ) as response:
+                        result = await response.json()
+                        # Validate response if security framework available
+                        if SECURITY_FRAMEWORK_AVAILABLE and self.security_manager:
+                            self.security_manager.validate_server_response(
+                                dict(response.headers)
+                            )
+                        if response.status >= 400:
+                            print(
+                                f"Request failed with status {response.status}: {result}"
+                            )
+                            return {"error": result, "status": response.status}
+                        return result
+                except Exception as e:
+                    print(f"Request attempt {attempt + 1} failed: {e}")
+                    if attempt < self.retry_attempts - 1:
+                        await asyncio.sleep(self.retry_delay)
+                    else:
+                        raise
+        except Exception as e:
+            print(f"Request failed: {e}")
+            raise
+
+    async def get(self, endpoint: str, **kwargs) -> Dict[str, Any]:
+        """Make GET request."""
+        return await self.request("GET", endpoint, **kwargs)
+
+    async def post(
+        self, endpoint: str, data: Dict[str, Any], **kwargs
+    ) -> Dict[str, Any]:
+        """Make POST request."""
+        return await self.request("POST", endpoint, data=data, **kwargs)
+
+    async def put(
+        self, endpoint: str, data: Dict[str, Any], **kwargs
+    ) -> Dict[str, Any]:
+        """Make PUT request."""
+        return await self.request("PUT", endpoint, data=data, **kwargs)
+
+    async def delete(self, endpoint: str, **kwargs) -> Dict[str, Any]:
+        """Make DELETE request."""
+        return await self.request("DELETE", endpoint, **kwargs)
+
+    async def test_connection(self) -> bool:
+        """Test connection to server."""
+        try:
+            result = await self.get("/health")
+            if "error" not in result:
+                print("✅ Connection test successful")
+                return True
+            else:
+                print(f"❌ Connection test failed: {result}")
+                return False
+        except Exception as e:
+            print(f"❌ Connection test failed: {e}")
+            return False
+
+    async def test_security_features(self) -> Dict[str, bool]:
+        """Test various security features."""
+        results = {}
+        # Test basic connectivity
+        results["connectivity"] = await self.test_connection()
+        # Test authentication
+        if self.auth_method != "none":
+            try:
+                result = await self.get("/api/auth/status")
+                results["authentication"] = "error" not in result
+            except:
+                results["authentication"] = False
+        # Test SSL/TLS
+        if self.base_url.startswith("https"):
+            results["ssl_tls"] = True
+        else:
+            results["ssl_tls"] = False
+        # Test certificate validation
+        if self.auth_method == "certificate" and SECURITY_FRAMEWORK_AVAILABLE:
+            results["certificate_validation"] = True
+        else:
+            results["certificate_validation"] = False
+        return results
+
+
+def create_client_config(
+    server_url: str, auth_method: str = "none", **kwargs
+) -> Dict[str, Any]:
+    """
+    Create client configuration for different authentication methods.
+    Args:
+        server_url: Server URL
+        auth_method: Authentication method (none, api_key, jwt, certificate, basic)
+        **kwargs: Additional configuration parameters
+    Returns:
+        Client configuration dictionary
+    """
+    config = {
+        "server_url": server_url,
+        "timeout": 30,
+        "retry_attempts": 3,
+        "retry_delay": 1,
+        "security": {"auth_method": auth_method},
+    }
+    if auth_method == "api_key":
+        config["security"]["api_key"] = {
+            "key": kwargs.get("api_key", "your_api_key_here"),
+            "header": kwargs.get("header", "X-API-Key"),
+        }
+    elif auth_method == "jwt":
+        config["security"]["jwt"] = {
+            "username": kwargs.get("username", "user"),
+            "password": kwargs.get("password", "password"),
+            "secret": kwargs.get("secret", "your_jwt_secret"),
+            "expiry_hours": kwargs.get("expiry_hours", 24),
+        }
+    elif auth_method == "certificate":
+        config["security"]["certificate"] = {
+            "enabled": True,
+            "cert_file": kwargs.get("cert_file", "./certs/client.crt"),
+            "key_file": kwargs.get("key_file", "./keys/client.key"),
+            "ca_cert_file": kwargs.get("ca_cert_file", "./certs/ca.crt"),
+        }
+        config["security"]["ssl"] = {
+            "enabled": True,
+            "check_hostname": kwargs.get("check_hostname", True),
+            "ca_cert_file": kwargs.get("ca_cert_file", "./certs/ca.crt"),
+        }
+    elif auth_method == "basic":
+        config["security"]["basic"] = {
+            "username": kwargs.get("username", "user"),
+            "password": kwargs.get("password", "password"),
+        }
+    return config
+
+
+async def demo_all_connection_methods():
+    """Demonstrate all possible connection methods."""
+    print("🚀 Universal Client Demo - All Connection Methods")
+    print("=" * 60)
+    # Test configurations for different auth methods
+    test_configs = [
+        {
+            "name": "No Authentication",
+            "config": create_client_config("http://localhost:8000", "none"),
+        },
+        {
+            "name": "API Key Authentication",
+            "config": create_client_config(
+                "http://localhost:8000", "api_key", api_key="demo_api_key_123"
+            ),
+        },
+        {
+            "name": "JWT Authentication",
+            "config": create_client_config(
+                "http://localhost:8000",
+                "jwt",
+                username="demo_user",
+                password="demo_password",
+                secret="demo_jwt_secret",
+            ),
+        },
+        {
+            "name": "Basic Authentication",
+            "config": create_client_config(
+                "http://localhost:8000",
+                "basic",
+                username="demo_user",
+                password="demo_password",
+            ),
+        },
+        {
+            "name": "Certificate Authentication (HTTPS)",
+            "config": create_client_config(
+                "https://localhost:8443",
+                "certificate",
+                cert_file="./certs/client.crt",
+                key_file="./keys/client.key",
+                ca_cert_file="./certs/ca.crt",
+            ),
+        },
+    ]
+    for test_config in test_configs:
+        print(f"\n📋 Testing: {test_config['name']}")
+        print("-" * 40)
+        try:
+            async with UniversalClient(test_config["config"]) as client:
+                # Test connection
+                success = await client.test_connection()
+                if success:
+                    # Test security features
+                    security_results = await client.test_security_features()
+                    print("Security Features:")
+                    for feature, status in security_results.items():
+                        status_icon = "✅" if status else "❌"
+                        print(f"  {status_icon} {feature}: {status}")
+                    # Make a test API call
+                    try:
+                        result = await client.get("/api/status")
+                        print(f"API Status: {result}")
+                    except Exception as e:
+                        print(f"API call failed: {e}")
+                else:
+                    print("❌ Connection failed")
+        except Exception as e:
+            print(f"❌ Test failed: {e}")
+    print("\n🎉 Demo completed!")
+
+
+async def demo_specific_connection(auth_method: str, **kwargs):
+    """
+    Demo specific connection method.
+    Args:
+        auth_method: Authentication method to test
+        **kwargs: Configuration parameters
+    """
+    print(f"🚀 Testing {auth_method} connection")
+    print("=" * 40)
+    config = create_client_config("http://localhost:8000", auth_method, **kwargs)
+    async with UniversalClient(config) as client:
+        # Test connection
+        success = await client.test_connection()
+        if success:
+            print("✅ Connection successful!")
+            # Make some API calls
+            try:
+                # Get server status
+                status = await client.get("/api/status")
+                print(f"Server Status: {status}")
+                # Test command execution
+                command_data = {
+                    "jsonrpc": "2.0",
+                    "method": "test_command",
+                    "params": {"message": "Hello from universal client!"},
+                    "id": 1,
+                }
+                result = await client.post("/api/jsonrpc", command_data)
+                print(f"Command Result: {result}")
+            except Exception as e:
+                print(f"API calls failed: {e}")
+        else:
+            print("❌ Connection failed")
+
+
+if __name__ == "__main__":
+    import sys
+    import argparse
+    import json
+
+    parser = argparse.ArgumentParser(
+        description="Universal Client for MCP Proxy Adapter"
+    )
+    parser.add_argument("--config", help="Path to configuration file")
+    parser.add_argument("--method", help="JSON-RPC method to call")
+    parser.add_argument("--params", help="JSON-RPC parameters (JSON string)")
+    parser.add_argument("--auth-method", help="Authentication method")
+    parser.add_argument("--server-url", help="Server URL")
+    args = parser.parse_args()
+    if args.config:
+        # Load configuration from file
+        try:
+            with open(args.config, "r") as f:
+                config_data = json.load(f)
+            # Extract server configuration
+            server_config = config_data.get("server", {})
+            host = server_config.get("host", "127.0.0.1")
+            port = server_config.get("port", 8000)
+            # Determine protocol
+            ssl_config = config_data.get("ssl", {})
+            ssl_enabled = ssl_config.get("enabled", False)
+            protocol = "https" if ssl_enabled else "http"
+            server_url = f"{protocol}://{host}:{port}"
+            print(f"🚀 Testing --config connection")
+            print("=" * 40)
+            print(f"Universal client initialized with auth method: --config")
+            print(f"Connecting to {server_url} with --config authentication...")
+            # Create client configuration
+            client_config = {
+                "server_url": server_url,
+                "timeout": 30,
+                "retry_attempts": 3,
+                "retry_delay": 1,
+                "security": {"auth_method": "none"},
+            }
+            # Add SSL configuration if needed
+            if ssl_enabled:
+                client_config["security"]["ssl"] = {
+                    "enabled": True,
+                    "check_hostname": False,
+                    "verify": False,
+                }
+                # Add CA certificate if available
+                ca_cert = ssl_config.get("ca_cert")
+                if ca_cert and os.path.exists(ca_cert):
+                    client_config["security"]["ssl"]["ca_cert_file"] = ca_cert
+
+            async def test_config_connection():
+                async with UniversalClient(client_config) as client:
+                    # Test connection
+                    success = await client.test_connection()
+                    if success:
+                        print("No authentication required")
+                        print("Connection established successfully")
+                        if args.method:
+                            # Execute JSON-RPC method
+                            params = {}
+                            if args.params:
+                                try:
+                                    params = json.loads(args.params)
+                                except json.JSONDecodeError:
+                                    print("❌ Invalid JSON parameters")
+                                    return
+                            command_data = {
+                                "jsonrpc": "2.0",
+                                "method": args.method,
+                                "params": params,
+                                "id": 1,
+                            }
+                            try:
+                                result = await client.post("/api/jsonrpc", command_data)
+                                print(
+                                    f"✅ Method '{args.method}' executed successfully:"
+                                )
+                                print(json.dumps(result, indent=2))
+                            except Exception as e:
+                                print(f"❌ Method execution failed: {e}")
+                        else:
+                            # Default to help command
+                            command_data = {
+                                "jsonrpc": "2.0",
+                                "method": "help",
+                                "params": {},
+                                "id": 1,
+                            }
+                            try:
+                                result = await client.post("/api/jsonrpc", command_data)
+                                print("✅ Help command executed successfully:")
+                                print(json.dumps(result, indent=2))
+                            except Exception as e:
+                                print(f"❌ Help command failed: {e}")
+                    else:
+                        print("❌ Connection failed")
+                    print("Connection closed")
+
+            asyncio.run(test_config_connection())
+        except FileNotFoundError:
+            print(f"❌ Configuration file not found: {args.config}")
+        except json.JSONDecodeError:
+            print(f"❌ Invalid JSON in configuration file: {args.config}")
+        except Exception as e:
+            print(f"❌ Error loading configuration: {e}")
+    elif len(sys.argv) > 1:
+        # Test specific auth method (legacy mode)
+        auth_method = sys.argv[1]
+        kwargs = {}
+        if auth_method == "api_key":
+            kwargs["api_key"] = "demo_key_123"
+        elif auth_method == "jwt":
+            kwargs.update(
+                {
+                    "username": "demo_user",
+                    "password": "demo_password",
+                    "secret": "demo_secret",
+                }
+            )
+        elif auth_method == "certificate":
+            kwargs.update(
+                {
+                    "cert_file": "./certs/client.crt",
+                    "key_file": "./keys/client.key",
+                    "ca_cert_file": "./certs/ca.crt",
+                }
+            )
+        elif auth_method == "basic":
+            kwargs.update({"username": "demo_user", "password": "demo_password"})
+        asyncio.run(demo_specific_connection(auth_method, **kwargs))
+    else:
+        # Demo all connection methods
+        asyncio.run(demo_all_connection_methods())