mcp-proxy-adapter 6.3.4__py3-none-any.whl → 6.3.5__py3-none-any.whl

mcp_proxy_adapter/core/client.py

@@ -24,10 +24,26 @@ from requests.exceptions import RequestException
 
 # Import security framework components
 try:
-    from mcp_security_framework import SecurityManager, AuthManager, CertificateManager, PermissionManager
-    from mcp_security_framework.utils import generate_api_key, create_jwt_token, validate_jwt_token
-    from mcp_security_framework.utils import extract_roles_from_cert, validate_certificate_chain
-    from mcp_security_framework.utils import create_ssl_context, validate_server_certificate
+    from mcp_security_framework import (
+        SecurityManager,
+        AuthManager,
+        CertificateManager,
+        PermissionManager,
+    )
+    from mcp_security_framework.utils import (
+        generate_api_key,
+        create_jwt_token,
+        validate_jwt_token,
+    )
+    from mcp_security_framework.utils import (
+        extract_roles_from_cert,
+        validate_certificate_chain,
+    )
+    from mcp_security_framework.utils import (
+        create_ssl_context,
+        validate_server_certificate,
+    )
+
     SECURITY_FRAMEWORK_AVAILABLE = True
 except ImportError:
     SECURITY_FRAMEWORK_AVAILABLE = False
@@ -36,7 +52,7 @@ except ImportError:
 class UniversalClient:
     """
     Universal client that demonstrates all possible secure connection methods.
-    
+
     Supports:
     - HTTP/HTTPS connections
     - API Key authentication
@@ -46,11 +62,11 @@ class UniversalClient:
     - 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
         """
@@ -59,72 +75,74 @@ class UniversalClient:
         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...")
-        
+        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()
@@ -136,110 +154,110 @@ class UniversalClient:
             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)
+                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)
+                "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")
@@ -248,9 +266,9 @@ class UniversalClient:
             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", {})
@@ -296,33 +314,33 @@ class UniversalClient:
         except Exception as e:
             print(f"Warning: Failed to create SSL context: {e}")
             return None
-    
+
     async def request(
-        self, 
-        method: str, 
-        endpoint: str, 
+        self,
+        method: str,
+        endpoint: str,
         data: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = 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:
@@ -331,20 +349,24 @@ class UniversalClient:
                         url,
                         json=data,
                         headers=request_headers,
-                        timeout=aiohttp.ClientTimeout(total=self.timeout)
+                        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))
-                        
+                            self.security_manager.validate_server_response(
+                                dict(response.headers)
+                            )
+
                         if response.status >= 400:
-                            print(f"Request failed with status {response.status}: {result}")
+                            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:
@@ -354,23 +376,27 @@ class UniversalClient:
         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]:
+
+    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]:
+
+    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:
@@ -384,14 +410,14 @@ class UniversalClient:
         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:
@@ -399,55 +425,58 @@ class UniversalClient:
                 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
-        
+
         return results
-    
+
     async def register_proxy(self, proxy_config: Dict[str, Any]) -> Dict[str, Any]:
         """
         Register with proxy server.
-        
+
         Args:
             proxy_config: Proxy registration configuration
-            
+
         Returns:
             Registration result
         """
         try:
-            result = await self.post("/api/jsonrpc", {
-                "jsonrpc": "2.0",
-                "method": "proxy_registration",
-                "params": proxy_config,
-                "id": 1
-            })
+            result = await self.post(
+                "/api/jsonrpc",
+                {
+                    "jsonrpc": "2.0",
+                    "method": "proxy_registration",
+                    "params": proxy_config,
+                    "id": 1,
+                },
+            )
             return result
         except Exception as e:
             print(f"Proxy registration failed: {e}")
             return {"error": str(e)}
-    
-    async def execute_command(self, command: str, params: Dict[str, Any] = None) -> Dict[str, Any]:
+
+    async def execute_command(
+        self, command: str, params: Dict[str, Any] = None
+    ) -> Dict[str, Any]:
         """
         Execute a command on the server.
-        
+
         Args:
             command: Command name
             params: Command parameters
-            
+
         Returns:
             Command result
         """
         try:
-            result = await self.post("/api/jsonrpc", {
-                "jsonrpc": "2.0",
-                "method": command,
-                "params": params or {},
-                "id": 1
-            })
+            result = await self.post(
+                "/api/jsonrpc",
+                {"jsonrpc": "2.0", "method": command, "params": params or {}, "id": 1},
+            )
             return result
         except Exception as e:
             print(f"Command execution failed: {e}")
@@ -457,55 +486,53 @@ class UniversalClient:
 def create_client_from_config(config_file: str) -> UniversalClient:
     """
     Create a UniversalClient instance from a configuration file.
-    
+
     Args:
         config_file: Path to configuration file
-        
+
     Returns:
         UniversalClient instance
     """
     try:
-        with open(config_file, 'r') as f:
+        with open(config_file, "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}"
-        
+
         # Create client configuration
         client_config = {
             "server_url": server_url,
             "timeout": 30,
             "retry_attempts": 3,
             "retry_delay": 1,
-            "security": {
-                "auth_method": "none"
-            }
+            "security": {"auth_method": "none"},
         }
-        
+
         # Add SSL configuration if needed
         if ssl_enabled:
             client_config["security"]["ssl"] = {
                 "enabled": True,
                 "check_hostname": False,
-                "verify": 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
-        
+
         return UniversalClient(client_config)
-        
+
     except Exception as e:
         raise ValueError(f"Failed to create client from config: {e}")
 
@@ -514,34 +541,36 @@ def create_client_from_config(config_file: str) -> UniversalClient:
 async def main():
     """Main function for CLI usage."""
     import argparse
-    
-    parser = argparse.ArgumentParser(description="Universal Client for MCP Proxy Adapter")
+
+    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:
             client = create_client_from_config(args.config)
-            
+
             print(f"🚀 Testing --config connection")
             print("=" * 40)
             print(f"Universal client initialized with auth method: --config")
             print(f"Connecting to {client.base_url} with --config authentication...")
-            
+
             async with 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 = {}
@@ -551,7 +580,7 @@ async def main():
                             except json.JSONDecodeError:
                                 print("❌ Invalid JSON parameters")
                                 return
-                        
+
                         result = await client.execute_command(args.method, params)
                         print(f"✅ Method '{args.method}' executed successfully:")
                         print(json.dumps(result, indent=2))
@@ -563,7 +592,7 @@ async def main():
                 else:
                     print("❌ Connection failed")
                 print("Connection closed")
-            
+
         except FileNotFoundError:
             print(f"❌ Configuration file not found: {args.config}")
         except json.JSONDecodeError: