SSL Proxy

system/proxy_module/proxy_manager.py

@@ -0,0 +1,235 @@
+import os
+import logging
+import subprocess
+import threading
+from pathlib import Path
+from typing import Optional, Dict, Any
+from dataclasses import asdict
+import psutil
+from logging.handlers import RotatingFileHandler
+import logging
+import os
+from .security import SecurityManager
+from .health_monitor import HealthMonitor
+from .cert_manager import CertificateManager
+from .metrics import MetricsCollector
+
+class ProxyManager:
+    def __init__(self, config_manager, base_path: str = None):
+        self.config_manager = config_manager
+        # Set up base paths first
+        self.base_path = Path(base_path) if base_path else Path(__file__).parent
+        
+        # Create required directories
+        self.certs_path = self.base_path / "certificates"
+        self.logs_path = self.base_path / "logs"
+        self.config_path = self.base_path / "config"
+        
+        for path in [self.certs_path, self.logs_path, self.config_path]:
+            path.mkdir(parents=True, exist_ok=True)
+        
+        # Setup logging before managers
+        self.setup_logging()
+        
+        # Initialize managers after paths are created
+        self.security = SecurityManager(self)
+        self.health_monitor = HealthMonitor(self)
+        self.cert_manager = CertificateManager(self)
+        self.metrics = MetricsCollector(self)
+        
+        self.proxy_process = None
+        self.logger.info("Proxy manager initialized successfully")
+
+    def setup_logging(self):
+        """Setup logging configuration with customized format"""
+        config = self.config_manager.get_instance()
+        log_level = getattr(logging, config.proxy_settings.log_level, logging.INFO)
+        
+        # Configure rotating file handler
+        log_file = self.logs_path / "proxy.log"
+        file_handler = RotatingFileHandler(
+            log_file,
+            maxBytes=5*1024*1024,
+            backupCount=5,
+            encoding='utf-8'
+        )
+        
+        # Configure logging formats
+        file_format = logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        console_format = logging.Formatter(
+            '[AllTalk PRX] %(message)s'
+        )
+        
+        file_handler.setFormatter(file_format)
+        
+        # Setup logger
+        self.logger = logging.getLogger("ProxyManager")
+        self.logger.setLevel(log_level)
+        self.logger.addHandler(file_handler)
+        
+        # Only add console handler if proxy debugging is enabled
+        if config.debugging.debug_proxy:
+            console_handler = logging.StreamHandler()
+            console_handler.setFormatter(console_format)
+            self.logger.addHandler(console_handler)
+            self.logger.info("Proxy debug logging enabled")
+
+    def _monitor_output(self, pipe, name):
+        """Monitor subprocess output and redirect to logger"""
+        try:
+            for line in iter(pipe.readline, ''):
+                line = line.strip()
+                if line:
+                    if name == "stderr":
+                        self.logger.error(line)
+                    else:
+                        # Pass through with our logging format
+                        print(line)  # Direct print for formatted output
+        except Exception as e:
+            self.logger.error(f"Error monitoring {name}: {e}")
+        finally:
+            pipe.close()
+
+    def start_proxy(self) -> bool:
+        """Start the proxy server with current configuration"""
+        if self.proxy_process and self.proxy_process.poll() is None:
+            self.logger.warning("Proxy already running")
+            return False
+
+        config = self.config_manager.get_instance()
+        proxy_settings = config.proxy_settings
+
+        try:
+
+            cmd = [
+                "python", "-m", "twisted_server",
+                "--external_ip_a", proxy_settings.api_endpoint.external_ip if proxy_settings.api_endpoint.enabled == "Enabled" else "0",
+                "--internal_ip_a", "127.0.0.1",
+                "--external_ip_b", proxy_settings.gradio_endpoint.external_ip if proxy_settings.gradio_endpoint.enabled == "Enabled" else "0",
+                "--internal_ip_b", "127.0.0.1",
+                "--cert_file_a", proxy_settings.api_endpoint.cert_name or "0",
+                "--cert_file_b", proxy_settings.gradio_endpoint.cert_name or "0"
+            ]            
+            
+            # Start process with pipe for output
+            self.proxy_process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                cwd=str(self.base_path),
+                text=True,
+                bufsize=1,
+                env={**os.environ, 'PYTHONUNBUFFERED': '1'},  # Force unbuffered output
+                universal_newlines=True
+            )
+            
+            # Start output monitoring threads
+            threading.Thread(target=self._monitor_output, 
+                           args=(self.proxy_process.stdout, "stdout"), 
+                           daemon=True).start()
+            threading.Thread(target=self._monitor_output, 
+                           args=(self.proxy_process.stderr, "stderr"), 
+                           daemon=True).start()
+            
+            if proxy_settings.proxy_enabled:
+                self.health_monitor.start_monitoring()
+                self.metrics.start_collecting()
+            
+            self.logger.info("Proxy server started successfully")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to start proxy: {e}")
+            return False
+
+    def stop_proxy(self) -> bool:
+        """Stop the proxy server and all monitoring"""
+        if not self.proxy_process:
+            self.logger.warning("No proxy process running")
+            return False
+
+        try:
+            # Stop all monitoring first
+            self.health_monitor.stop_monitoring()
+            self.metrics.stop_collecting()
+            
+            # Stop the main process
+            process = psutil.Process(self.proxy_process.pid)
+            for proc in process.children(recursive=True):
+                proc.terminate()
+            process.terminate()
+            
+            try:
+                process.wait(timeout=5)
+            except psutil.TimeoutExpired:
+                process.kill()
+            
+            self.proxy_process = None
+            self.logger.info("Proxy server stopped successfully")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to stop proxy: {e}")
+            return False
+
+    def get_status(self) -> Dict[str, Any]:
+        """Get comprehensive status of the proxy system"""
+        config = self.config_manager.get_instance()
+        proxy_config = config.proxy_settings
+        
+        status = {
+            "running": self.proxy_process is not None and self.proxy_process.poll() is None,
+            "config": {
+                "gradio_endpoint": asdict(proxy_config.gradio_endpoint),
+                "api_endpoint": asdict(proxy_config.api_endpoint),
+                "start_on_startup": proxy_config.start_on_startup
+            },
+            "certificates": self.cert_manager.get_certificates_status(),
+            "health": self.health_monitor.get_health_status(),
+            "metrics": self.metrics.get_current_metrics(),
+            "security": self.security.get_security_status()
+        }
+        
+        return status
+
+    def handle_cert_upload(self, cert_file, key_file, name):
+        """Handle certificate upload from the interface"""
+        if not cert_file or not key_file or not name:
+            return "Please provide both certificate and key files, and a name"
+        
+        cert_path = Path(cert_file.name)
+        key_path = Path(key_file.name)
+        
+        success = self.cert_manager.install_certificate(cert_path, key_path, name)
+        return "Certificate uploaded successfully" if success else "Failed to upload certificate"
+
+    def handle_start(self):
+        """Handle proxy start from the interface"""
+        success = self.start_proxy()
+        return "Proxy Server: Starting..." if success else "Failed to start proxy server"
+
+    def handle_stop(self):
+        """Handle proxy stop from the interface"""
+        success = self.stop_proxy()
+        return "Proxy Server: Stopped" if success else "Failed to stop proxy server"
+
+    def handle_status(self):
+        """Handle status request from the interface"""
+        status = self.get_status()
+        if status['running']:
+            return "Proxy Server: Running\n" + \
+                f"API Endpoint: {'Active' if status['config']['api_endpoint']['enabled'] == 'Enabled' else 'Disabled'}\n" + \
+                f"Gradio Endpoint: {'Active' if status['config']['gradio_endpoint']['enabled'] == 'Enabled' else 'Disabled'}\n" + \
+                f"Certificates: {len(status['certificates'])} configured"
+        else:
+            return "Proxy Server: Stopped"
+
+    def handle_security_event(self, event_type: str, details: Dict[str, Any]):
+        """Handle security events from the SecurityManager"""
+        self.logger.warning(f"Security event: {event_type} - {details}")
+        # Implement security response logic here
+
+    def cleanup(self):
+        """Cleanup resources before shutdown"""
+        self.stop_proxy()
+        self.logger.info("Proxy manager cleanup completed")