mcp-proxy-adapter 6.6.9__py3-none-any.whl → 6.7.1__py3-none-any.whl

mcp_proxy_adapter/core/mtls_proxy.py

@@ -0,0 +1,181 @@
+"""
+mTLS Proxy for MCP Proxy Adapter
+
+This module provides mTLS proxy functionality that accepts mTLS connections
+and proxies them to the internal hypercorn server.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import asyncio
+import ssl
+import logging
+from typing import Optional, Dict, Any
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class MTLSProxy:
+    """
+    mTLS Proxy that accepts mTLS connections and proxies them to internal server.
+    """
+    
+    def __init__(self, 
+                 external_host: str,
+                 external_port: int,
+                 internal_host: str = "127.0.0.1",
+                 internal_port: int = 9000,
+                 cert_file: Optional[str] = None,
+                 key_file: Optional[str] = None,
+                 ca_cert: Optional[str] = None):
+        """
+        Initialize mTLS Proxy.
+        
+        Args:
+            external_host: External host to bind to
+            external_port: External port to bind to
+            internal_host: Internal server host
+            internal_port: Internal server port
+            cert_file: Server certificate file
+            key_file: Server private key file
+            ca_cert: CA certificate file for client verification
+        """
+        self.external_host = external_host
+        self.external_port = external_port
+        self.internal_host = internal_host
+        self.internal_port = internal_port
+        self.cert_file = cert_file
+        self.key_file = key_file
+        self.ca_cert = ca_cert
+        self.server = None
+        
+    async def start(self):
+        """Start the mTLS proxy server."""
+        try:
+            # Create SSL context
+            ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+            ssl_context.load_cert_chain(self.cert_file, self.key_file)
+            
+            if self.ca_cert:
+                ssl_context.load_verify_locations(self.ca_cert)
+                ssl_context.verify_mode = ssl.CERT_REQUIRED
+            else:
+                ssl_context.verify_mode = ssl.CERT_NONE
+                
+            # Start server
+            self.server = await asyncio.start_server(
+                self._handle_client,
+                self.external_host,
+                self.external_port,
+                ssl=ssl_context
+            )
+            
+            logger.info(f"🔐 mTLS Proxy started on {self.external_host}:{self.external_port}")
+            logger.info(f"🌐 Proxying to {self.internal_host}:{self.internal_port}")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to start mTLS proxy: {e}")
+            raise
+            
+    async def stop(self):
+        """Stop the mTLS proxy server."""
+        if self.server:
+            self.server.close()
+            await self.server.wait_closed()
+            logger.info("🔐 mTLS Proxy stopped")
+            
+    async def _handle_client(self, reader, writer):
+        """Handle client connection."""
+        try:
+            # Get client address
+            client_addr = writer.get_extra_info('peername')
+            logger.info(f"🔐 mTLS connection from {client_addr}")
+            
+            # Connect to internal server
+            internal_reader, internal_writer = await asyncio.open_connection(
+                self.internal_host, self.internal_port
+            )
+            
+            # Create bidirectional proxy
+            await asyncio.gather(
+                self._proxy_data(reader, internal_writer, "client->server"),
+                self._proxy_data(internal_reader, writer, "server->client")
+            )
+            
+        except Exception as e:
+            logger.error(f"❌ Error handling client connection: {e}")
+        finally:
+            try:
+                writer.close()
+                await writer.wait_closed()
+            except:
+                pass
+                
+    async def _proxy_data(self, reader, writer, direction):
+        """Proxy data between reader and writer."""
+        try:
+            while True:
+                data = await reader.read(4096)
+                if not data:
+                    break
+                writer.write(data)
+                await writer.drain()
+        except Exception as e:
+            logger.debug(f"Proxy connection closed ({direction}): {e}")
+        finally:
+            try:
+                writer.close()
+                await writer.wait_closed()
+            except:
+                pass
+
+
+async def start_mtls_proxy(config: Dict[str, Any]) -> Optional[MTLSProxy]:
+    """
+    Start mTLS proxy based on configuration.
+    
+    Args:
+        config: Application configuration
+        
+    Returns:
+        MTLSProxy instance if started, None otherwise
+    """
+    # Check if mTLS is enabled
+    protocol = config.get("server", {}).get("protocol", "http")
+    verify_client = config.get("transport", {}).get("verify_client", False)
+    
+    # Only start mTLS proxy if mTLS is explicitly enabled
+    if protocol != "mtls" and not verify_client:
+        logger.info("🌐 Regular mode: no mTLS proxy needed")
+        return None
+        
+    # Get configuration
+    server_config = config.get("server", {})
+    transport_config = config.get("transport", {})
+    
+    external_host = server_config.get("host", "0.0.0.0")
+    external_port = server_config.get("port", 8000)
+    internal_port = external_port + 1000  # Internal port
+    
+    cert_file = transport_config.get("cert_file")
+    key_file = transport_config.get("key_file")
+    ca_cert = transport_config.get("ca_cert")
+    
+    if not cert_file or not key_file:
+        logger.warning("⚠️  mTLS enabled but certificates not configured")
+        return None
+        
+    # Create and start proxy
+    proxy = MTLSProxy(
+        external_host=external_host,
+        external_port=external_port,
+        internal_port=internal_port,
+        cert_file=cert_file,
+        key_file=key_file,
+        ca_cert=ca_cert
+    )
+    
+    await proxy.start()
+    return proxy

mcp_proxy_adapter/core/signal_handler.py

@@ -0,0 +1,170 @@
+"""
+Signal handler for graceful shutdown with proxy unregistration.
+
+This module provides signal handling for SIGTERM, SIGINT, and SIGHUP
+to ensure proper proxy unregistration before server shutdown.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import signal
+import asyncio
+import threading
+from typing import Optional, Callable, Any
+from mcp_proxy_adapter.core.logging import logger
+
+
+class SignalHandler:
+    """
+    Signal handler for graceful shutdown with proxy unregistration.
+    """
+    
+    def __init__(self):
+        """Initialize signal handler."""
+        self._shutdown_callback: Optional[Callable] = None
+        self._shutdown_event = threading.Event()
+        self._original_handlers = {}
+        self._setup_signal_handlers()
+    
+    def set_shutdown_callback(self, callback: Callable):
+        """
+        Set callback function to be called during shutdown.
+        
+        Args:
+            callback: Function to call during shutdown
+        """
+        self._shutdown_callback = callback
+        logger.info("Shutdown callback set for signal handler")
+    
+    def _setup_signal_handlers(self):
+        """Setup signal handlers for graceful shutdown."""
+        # Handle SIGTERM (termination signal)
+        self._original_handlers[signal.SIGTERM] = signal.signal(
+            signal.SIGTERM, self._handle_shutdown_signal
+        )
+        
+        # Handle SIGINT (Ctrl+C)
+        self._original_handlers[signal.SIGINT] = signal.signal(
+            signal.SIGINT, self._handle_shutdown_signal
+        )
+        
+        # Handle SIGHUP (hangup signal)
+        self._original_handlers[signal.SIGHUP] = signal.signal(
+            signal.SIGHUP, self._handle_shutdown_signal
+        )
+        
+        logger.info("Signal handlers installed for SIGTERM, SIGINT, SIGHUP")
+    
+    def _handle_shutdown_signal(self, signum: int, frame: Any):
+        """
+        Handle shutdown signals.
+        
+        Args:
+            signum: Signal number
+            frame: Current stack frame
+        """
+        signal_name = signal.Signals(signum).name
+        logger.info(f"🛑 Received {signal_name} signal, initiating graceful shutdown...")
+        
+        # Set shutdown event
+        self._shutdown_event.set()
+        
+        # Call shutdown callback if set
+        if self._shutdown_callback:
+            try:
+                logger.info("🔄 Executing shutdown callback...")
+                self._shutdown_callback()
+                logger.info("✅ Shutdown callback completed successfully")
+            except Exception as e:
+                logger.error(f"❌ Shutdown callback failed: {e}")
+        
+        # If this is SIGINT (Ctrl+C), we might want to exit immediately after a delay
+        if signum == signal.SIGINT:
+            logger.info("⏰ SIGINT received, will exit in 5 seconds if not stopped gracefully...")
+            threading.Timer(5.0, self._force_exit).start()
+    
+    def _force_exit(self):
+        """Force exit if graceful shutdown takes too long."""
+        if self._shutdown_event.is_set():
+            logger.warning("⚠️ Forcing exit after timeout")
+            import os
+            os._exit(1)
+    
+    def wait_for_shutdown(self, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for shutdown signal.
+        
+        Args:
+            timeout: Maximum time to wait in seconds
+            
+        Returns:
+            True if shutdown signal received, False if timeout
+        """
+        return self._shutdown_event.wait(timeout)
+    
+    def is_shutdown_requested(self) -> bool:
+        """
+        Check if shutdown has been requested.
+        
+        Returns:
+            True if shutdown signal received
+        """
+        return self._shutdown_event.is_set()
+    
+    def restore_handlers(self):
+        """Restore original signal handlers."""
+        for sig, handler in self._original_handlers.items():
+            if handler is not None:
+                signal.signal(sig, handler)
+        logger.info("Original signal handlers restored")
+
+
+# Global signal handler instance
+_signal_handler: Optional[SignalHandler] = None
+
+
+def get_signal_handler() -> SignalHandler:
+    """Get the global signal handler instance."""
+    global _signal_handler
+    if _signal_handler is None:
+        _signal_handler = SignalHandler()
+    return _signal_handler
+
+
+def setup_signal_handling(shutdown_callback: Optional[Callable] = None):
+    """
+    Setup signal handling for graceful shutdown.
+    
+    Args:
+        shutdown_callback: Optional callback to execute during shutdown
+    """
+    handler = get_signal_handler()
+    if shutdown_callback:
+        handler.set_shutdown_callback(shutdown_callback)
+    logger.info("Signal handling setup completed")
+
+
+def wait_for_shutdown_signal(timeout: Optional[float] = None) -> bool:
+    """
+    Wait for shutdown signal.
+    
+    Args:
+        timeout: Maximum time to wait in seconds
+        
+    Returns:
+        True if shutdown signal received, False if timeout
+    """
+    handler = get_signal_handler()
+    return handler.wait_for_shutdown(timeout)
+
+
+def is_shutdown_requested() -> bool:
+    """
+    Check if shutdown has been requested.
+    
+    Returns:
+        True if shutdown signal received
+    """
+    handler = get_signal_handler()
+    return handler.is_shutdown_requested()