mcp-proxy-adapter 6.0.0__py3-none-any.whl → 6.1.1__py3-none-any.whl

mcp_proxy_adapter/core/proxy_registration.py

@@ -1,13 +1,18 @@
 """
-Module for proxy registration functionality.
+Module for proxy registration functionality with security framework integration.
 
 This module handles automatic registration and unregistration of the server
-with the MCP proxy server during startup and shutdown.
+with the MCP proxy server during startup and shutdown, using mcp_security_framework
+for secure connections and authentication.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
 """
 
 import asyncio
 import json
 import time
+import ssl
 from typing import Dict, Any, Optional, Tuple
 from urllib.parse import urljoin
 
@@ -15,8 +20,8 @@ import aiohttp
 import requests
 from requests.exceptions import RequestException
 
-from mcp_proxy_adapter.config import config
 from mcp_proxy_adapter.core.logging import logger
+from mcp_proxy_adapter.core.client_security import create_client_security_manager
 
 
 class ProxyRegistrationError(Exception):
@@ -26,28 +31,49 @@ class ProxyRegistrationError(Exception):
 
 class ProxyRegistrationManager:
     """
-    Manager for proxy registration functionality.
+    Manager for proxy registration functionality with security framework integration.
     
     Handles automatic registration and unregistration of the server
-    with the MCP proxy server.
+    with the MCP proxy server using secure authentication methods.
     """
     
-    def __init__(self):
-        """Initialize the proxy registration manager."""
-        self.registration_config = config.get("proxy_registration", {})
-        self.proxy_url = self.registration_config.get("proxy_url", "http://localhost:3004")
-        self.server_id = self.registration_config.get("server_id", "mcp_proxy_adapter")
-        self.server_name = self.registration_config.get("server_name", "MCP Proxy Adapter")
-        self.description = self.registration_config.get("description", "JSON-RPC API for interacting with MCP Proxy")
-        self.timeout = self.registration_config.get("registration_timeout", 30)
-        self.retry_attempts = self.registration_config.get("retry_attempts", 3)
-        self.retry_delay = self.registration_config.get("retry_delay", 5)
-        self.auto_register = self.registration_config.get("auto_register_on_startup", True)
-        self.auto_unregister = self.registration_config.get("auto_unregister_on_shutdown", True)
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initialize the proxy registration manager.
+        
+        Args:
+            config: Application configuration
+        """
+        self.config = config
+        self.registration_config = config.get("registration", {})
         
+        # Basic registration settings
+        self.proxy_url = self.registration_config.get("server_url", "https://proxy-registry.example.com")
+        self.server_id = self.registration_config.get("proxy_info", {}).get("name", "mcp_proxy_adapter")
+        self.server_name = self.registration_config.get("proxy_info", {}).get("name", "MCP Proxy Adapter")
+        self.description = self.registration_config.get("proxy_info", {}).get("description", "JSON-RPC API for interacting with MCP Proxy")
+        
+        # Heartbeat settings
+        heartbeat_config = self.registration_config.get("heartbeat", {})
+        self.timeout = heartbeat_config.get("timeout", 30)
+        self.retry_attempts = heartbeat_config.get("retry_attempts", 3)
+        self.retry_delay = heartbeat_config.get("retry_delay", 60)
+        self.heartbeat_interval = heartbeat_config.get("interval", 300)
+        
+        # Auto registration settings
+        self.auto_register = self.registration_config.get("enabled", False)
+        self.auto_unregister = True  # Always unregister on shutdown
+        
+        # Initialize client security manager
+        self.client_security = create_client_security_manager(config)
+        
+        # Registration state
         self.registered = False
         self.server_key: Optional[str] = None
         self.server_url: Optional[str] = None
+        self.heartbeat_task: Optional[asyncio.Task] = None
+        
+        logger.info("Proxy registration manager initialized with security framework integration")
         
     def is_enabled(self) -> bool:
         """
@@ -68,9 +94,55 @@ class ProxyRegistrationManager:
         self.server_url = server_url
         logger.info(f"Proxy registration server URL set to: {server_url}")
     
+    def _get_auth_headers(self) -> Dict[str, str]:
+        """
+        Get authentication headers for registration requests.
+        
+        Returns:
+            Dictionary of authentication headers
+        """
+        if not self.client_security:
+            return {"Content-Type": "application/json"}
+        
+        auth_method = self.registration_config.get("auth_method", "certificate")
+        
+        if auth_method == "certificate":
+            return self.client_security.get_client_auth_headers("certificate")
+        elif auth_method == "token":
+            token_config = self.registration_config.get("token", {})
+            token = token_config.get("token")
+            return self.client_security.get_client_auth_headers("jwt", token=token)
+        elif auth_method == "api_key":
+            api_key_config = self.registration_config.get("api_key", {})
+            api_key = api_key_config.get("key")
+            return self.client_security.get_client_auth_headers("api_key", api_key=api_key)
+        else:
+            return {"Content-Type": "application/json"}
+    
+    def _create_ssl_context(self) -> Optional[ssl.SSLContext]:
+        """
+        Create SSL context for secure connections.
+        
+        Returns:
+            SSL context or None if SSL not needed
+        """
+        if not self.client_security:
+            return None
+        
+        try:
+            # Check if SSL is enabled for registration
+            cert_config = self.registration_config.get("certificate", {})
+            if cert_config.get("enabled", False):
+                return self.client_security.create_client_ssl_context()
+            
+            return None
+        except Exception as e:
+            logger.error(f"Failed to create SSL context: {e}")
+            return None
+    
     async def register_server(self) -> bool:
         """
-        Register the server with the proxy.
+        Register the server with the proxy using secure authentication.
         
         Returns:
             True if registration was successful, False otherwise.
@@ -83,11 +155,20 @@ class ProxyRegistrationManager:
             logger.error("Server URL not set, cannot register with proxy")
             return False
         
+        # Prepare registration data with proxy info
+        proxy_info = self.registration_config.get("proxy_info", {})
         registration_data = {
             "server_id": self.server_id,
             "server_url": self.server_url,
             "server_name": self.server_name,
-            "description": self.description
+            "description": self.description,
+            "version": proxy_info.get("version", "1.0.0"),
+            "capabilities": proxy_info.get("capabilities", ["jsonrpc", "rest"]),
+            "endpoints": proxy_info.get("endpoints", {
+                "jsonrpc": "/api/jsonrpc",
+                "rest": "/cmd",
+                "health": "/health"
+            })
         }
         
         logger.info(f"Attempting to register server with proxy at {self.proxy_url}")
@@ -95,12 +176,17 @@ class ProxyRegistrationManager:
         
         for attempt in range(self.retry_attempts):
             try:
-                success, result = await self._make_registration_request(registration_data)
+                success, result = await self._make_secure_registration_request(registration_data)
                 
                 if success:
                     self.registered = True
                     self.server_key = result.get("server_key")
                     logger.info(f"✅ Successfully registered with proxy. Server key: {self.server_key}")
+                    
+                    # Start heartbeat if enabled
+                    if self.registration_config.get("heartbeat", {}).get("enabled", True):
+                        await self._start_heartbeat()
+                    
                     return True
                 else:
                     error_msg = result.get("error", {}).get("message", "Unknown error")
@@ -135,6 +221,9 @@ class ProxyRegistrationManager:
             logger.info("Server not registered with proxy, skipping unregistration")
             return True
         
+        # Stop heartbeat
+        await self._stop_heartbeat()
+        
         # Extract copy_number from server_key (format: server_id_copy_number)
         try:
             copy_number = int(self.server_key.split("_")[-1])
@@ -150,7 +239,7 @@ class ProxyRegistrationManager:
         logger.debug(f"Unregistration data: {unregistration_data}")
         
         try:
-            success, result = await self._make_unregistration_request(unregistration_data)
+            success, result = await self._make_secure_unregistration_request(unregistration_data)
             
             if success:
                 unregistered = result.get("unregistered", False)
@@ -171,9 +260,9 @@ class ProxyRegistrationManager:
             logger.error(f"❌ Unregistration failed with exception: {e}")
             return False
     
-    async def _make_registration_request(self, data: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
+    async def _make_secure_registration_request(self, data: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
         """
-        Make registration request to proxy.
+        Make secure registration request to proxy using security framework.
         
         Args:
             data: Registration data.
@@ -183,19 +272,40 @@ class ProxyRegistrationManager:
         """
         url = urljoin(self.proxy_url, "/register")
         
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                url,
-                json=data,
-                headers={"Content-Type": "application/json"},
-                timeout=aiohttp.ClientTimeout(total=self.timeout)
-            ) as response:
-                result = await response.json()
-                return response.status == 200, result
+        # Get authentication headers
+        headers = self._get_auth_headers()
+        headers["Content-Type"] = "application/json"
+        
+        # Create SSL context if needed
+        ssl_context = self._create_ssl_context()
+        
+        # Create connector with SSL context
+        connector = None
+        if ssl_context:
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+        
+        try:
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.post(
+                    url,
+                    json=data,
+                    headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=self.timeout)
+                ) as response:
+                    result = await response.json()
+                    
+                    # Validate response headers if security framework available
+                    if self.client_security:
+                        self.client_security.validate_server_response(dict(response.headers))
+                    
+                    return response.status == 200, result
+        finally:
+            if connector:
+                await connector.close()
     
-    async def _make_unregistration_request(self, data: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
+    async def _make_secure_unregistration_request(self, data: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
         """
-        Make unregistration request to proxy.
+        Make secure unregistration request to proxy using security framework.
         
         Args:
             data: Unregistration data.
@@ -205,15 +315,120 @@ class ProxyRegistrationManager:
         """
         url = urljoin(self.proxy_url, "/unregister")
         
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                url,
-                json=data,
-                headers={"Content-Type": "application/json"},
-                timeout=aiohttp.ClientTimeout(total=self.timeout)
-            ) as response:
-                result = await response.json()
-                return response.status == 200, result
+        # Get authentication headers
+        headers = self._get_auth_headers()
+        headers["Content-Type"] = "application/json"
+        
+        # Create SSL context if needed
+        ssl_context = self._create_ssl_context()
+        
+        # Create connector with SSL context
+        connector = None
+        if ssl_context:
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+        
+        try:
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.post(
+                    url,
+                    json=data,
+                    headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=self.timeout)
+                ) as response:
+                    result = await response.json()
+                    
+                    # Validate response headers if security framework available
+                    if self.client_security:
+                        self.client_security.validate_server_response(dict(response.headers))
+                    
+                    return response.status == 200, result
+        finally:
+            if connector:
+                await connector.close()
+    
+    async def _start_heartbeat(self) -> None:
+        """Start heartbeat task for keeping registration alive."""
+        if self.heartbeat_task and not self.heartbeat_task.done():
+            return
+        
+        self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+        logger.info("Heartbeat task started")
+    
+    async def _stop_heartbeat(self) -> None:
+        """Stop heartbeat task."""
+        if self.heartbeat_task and not self.heartbeat_task.done():
+            self.heartbeat_task.cancel()
+            try:
+                await self.heartbeat_task
+            except asyncio.CancelledError:
+                pass
+            logger.info("Heartbeat task stopped")
+    
+    async def _heartbeat_loop(self) -> None:
+        """Heartbeat loop to keep registration alive."""
+        while self.registered:
+            try:
+                await asyncio.sleep(self.heartbeat_interval)
+                
+                if not self.registered:
+                    break
+                
+                # Send heartbeat
+                success = await self._send_heartbeat()
+                if not success:
+                    logger.warning("Heartbeat failed, attempting to re-register")
+                    await self.register_server()
+                    
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Heartbeat error: {e}")
+    
+    async def _send_heartbeat(self) -> bool:
+        """Send heartbeat to proxy server."""
+        if not self.server_key:
+            return False
+        
+        heartbeat_data = {
+            "server_id": self.server_id,
+            "server_key": self.server_key,
+            "timestamp": int(time.time())
+        }
+        
+        url = urljoin(self.proxy_url, "/heartbeat")
+        
+        # Get authentication headers
+        headers = self._get_auth_headers()
+        headers["Content-Type"] = "application/json"
+        
+        # Create SSL context if needed
+        ssl_context = self._create_ssl_context()
+        
+        # Create connector with SSL context
+        connector = None
+        if ssl_context:
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+        
+        try:
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.post(
+                    url,
+                    json=heartbeat_data,
+                    headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=self.timeout)
+                ) as response:
+                    if response.status == 200:
+                        logger.debug("Heartbeat sent successfully")
+                        return True
+                    else:
+                        logger.warning(f"Heartbeat failed with status: {response.status}")
+                        return False
+        except Exception as e:
+            logger.error(f"Heartbeat error: {e}")
+            return False
+        finally:
+            if connector:
+                await connector.close()
     
     def get_registration_status(self) -> Dict[str, Any]:
         """
@@ -222,18 +437,44 @@ class ProxyRegistrationManager:
         Returns:
             Dictionary with registration status information.
         """
-        return {
+        status = {
             "enabled": self.is_enabled(),
             "registered": self.registered,
             "server_key": self.server_key,
             "server_url": self.server_url,
             "proxy_url": self.proxy_url,
-            "server_id": self.server_id
+            "server_id": self.server_id,
+            "heartbeat_active": self.heartbeat_task is not None and not self.heartbeat_task.done()
         }
+        
+        # Add security information if available
+        if self.client_security:
+            status["security_enabled"] = True
+            status["ssl_enabled"] = self.client_security.is_ssl_enabled()
+            status["auth_methods"] = self.client_security.get_supported_auth_methods()
+            
+            cert_info = self.client_security.get_client_certificate_info()
+            if cert_info:
+                status["client_certificate"] = cert_info
+        else:
+            status["security_enabled"] = False
+        
+        return status
+
 
+# Global proxy registration manager instance (will be initialized with config)
+proxy_registration_manager: Optional[ProxyRegistrationManager] = None
 
-# Global proxy registration manager instance
-proxy_registration_manager = ProxyRegistrationManager()
+
+def initialize_proxy_registration(config: Dict[str, Any]) -> None:
+    """
+    Initialize global proxy registration manager.
+    
+    Args:
+        config: Application configuration
+    """
+    global proxy_registration_manager
+    proxy_registration_manager = ProxyRegistrationManager(config)
 
 
 async def register_with_proxy(server_url: str) -> bool:
@@ -246,6 +487,10 @@ async def register_with_proxy(server_url: str) -> bool:
     Returns:
         True if registration was successful, False otherwise.
     """
+    if not proxy_registration_manager:
+        logger.error("Proxy registration manager not initialized")
+        return False
+    
     proxy_registration_manager.set_server_url(server_url)
     return await proxy_registration_manager.register_server()
 
@@ -257,6 +502,10 @@ async def unregister_from_proxy() -> bool:
     Returns:
         True if unregistration was successful, False otherwise.
     """
+    if not proxy_registration_manager:
+        logger.error("Proxy registration manager not initialized")
+        return False
+    
     return await proxy_registration_manager.unregister_server()
 
 
@@ -267,4 +516,7 @@ def get_proxy_registration_status() -> Dict[str, Any]:
     Returns:
         Dictionary with registration status information.
     """
+    if not proxy_registration_manager:
+        return {"error": "Proxy registration manager not initialized"}
+    
     return proxy_registration_manager.get_registration_status() 

mcp_proxy_adapter/core/security_adapter.py

@@ -16,8 +16,7 @@ try:
     from mcp_security_framework.schemas.config import (
         AuthConfig, SSLConfig, PermissionConfig, RateLimitConfig
     )
-    from mcp_security_framework.schemas.request import SecurityRequest
-    from mcp_security_framework.schemas.response import SecurityResult
+    # Note: SecurityRequest and SecurityResult are not available in current version
     SECURITY_FRAMEWORK_AVAILABLE = True
 except ImportError:
     # Fallback for when mcp_security_framework is not available
@@ -28,8 +27,6 @@ except ImportError:
     SSLConfig = None
     PermissionConfig = None
     RateLimitConfig = None
-    SecurityRequest = None
-    SecurityResult = None
 
 from mcp_proxy_adapter.core.logging import logger
 
@@ -245,24 +242,24 @@ class SecurityAdapter:
                 "user_id": None
             }
     
-    def _create_security_request(self, request_data: Dict[str, Any]) -> SecurityRequest:
+    def _create_security_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Create SecurityRequest from request data.
+        Create request data for security validation.
         
         Args:
             request_data: Request data dictionary
             
         Returns:
-            SecurityRequest instance
+            Request data dictionary for security validation
         """
-        return SecurityRequest(
-            method=request_data.get("method", "GET"),
-            path=request_data.get("path", "/"),
-            headers=request_data.get("headers", {}),
-            query_params=request_data.get("query_params", {}),
-            client_ip=request_data.get("client_ip", "unknown"),
-            body=request_data.get("body", {})
-        )
+        return {
+            "method": request_data.get("method", "GET"),
+            "path": request_data.get("path", "/"),
+            "headers": request_data.get("headers", {}),
+            "query_params": request_data.get("query_params", {}),
+            "client_ip": request_data.get("client_ip", "unknown"),
+            "body": request_data.get("body", {})
+        }
     
     def _fallback_validation(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
         """