diff --git a/system/proxy_module/cert_maker.py b/system/proxy_module/cert_maker.py new file mode 100644 index 0000000..46de06e --- /dev/null +++ b/system/proxy_module/cert_maker.py @@ -0,0 +1,51 @@ +from pathlib import Path +import argparse +from OpenSSL import crypto +import sys + +def generate_certificate(cert_name: str, output_dir: Path, cn: str = "localhost", days: int = 365): + k = crypto.PKey() + k.generate_key(crypto.TYPE_RSA, 2048) + + cert = crypto.X509() + cert.get_subject().CN = cn + cert.set_serial_number(1000) + cert.gmtime_adj_notBefore(0) + cert.gmtime_adj_notAfter(days*24*60*60) + cert.set_issuer(cert.get_subject()) + cert.set_pubkey(k) + cert.sign(k, 'sha256') + + output_dir.mkdir(parents=True, exist_ok=True) + + with open(output_dir / f"{cert_name}_cert.pem", "wb") as f: + f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert)) + with open(output_dir / f"{cert_name}_key.pem", "wb") as f: + f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k)) + +def main(): + print("Certificate Generator") + print("-" * 20) + + # Get default paths + default_path = Path(__file__).parent / "certificates" + + # Get user input with defaults + cert_name = input(f"Certificate name (default: test_cert): ") or "test_cert" + cn = input(f"Common Name (default: localhost): ") or "localhost" + days = input(f"Valid days (default: 365): ") or "365" + output_dir = input(f"Output directory (default: {default_path}): ") or str(default_path) + + try: + days = int(days) + output_path = Path(output_dir) + generate_certificate(cert_name, output_path, cn, days) + print(f"\nSuccess! Certificate files created:") + print(f"- {output_path / f'{cert_name}_cert.pem'}") + print(f"- {output_path / f'{cert_name}_key.pem'}") + except Exception as e: + print(f"\nError: {str(e)}") + sys.exit(1) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/system/proxy_module/cert_manager.py b/system/proxy_module/cert_manager.py new file mode 100644 index 0000000..c0b06e8 --- /dev/null +++ b/system/proxy_module/cert_manager.py @@ -0,0 +1,120 @@ +import logging +import os +from pathlib import Path +from typing import Dict, Any, Optional +from datetime import datetime +import ssl +import OpenSSL +from dataclasses import dataclass + +@dataclass +class CertificateInfo: + subject: str + issuer: str + valid_from: datetime + valid_until: datetime + serial_number: str + is_valid: bool + +class CertificateManager: + def __init__(self, proxy_manager): + self.proxy_manager = proxy_manager + self.logger = logging.getLogger("ProxyCertManager") + self.certs_path = proxy_manager.certs_path + + def validate_certificate(self, cert_path: Path) -> Optional[CertificateInfo]: + """Validate and get information about a certificate""" + try: + with open(cert_path, 'rb') as cert_file: + cert_data = cert_file.read() + x509 = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, cert_data) + + valid_from = datetime.strptime(x509.get_notBefore().decode(), '%Y%m%d%H%M%SZ') + valid_until = datetime.strptime(x509.get_notAfter().decode(), '%Y%m%d%H%M%SZ') + now = datetime.now() + + return CertificateInfo( + subject=x509.get_subject().CN, + issuer=x509.get_issuer().CN, + valid_from=valid_from, + valid_until=valid_until, + serial_number=str(x509.get_serial_number()), + is_valid=valid_from <= now <= valid_until + ) + except Exception as e: + self.logger.error(f"Certificate validation error: {e}") + return None + + def install_certificate(self, cert_path: Path, key_path: Path, name: str) -> bool: + try: + import shutil + # Validate certificate first + cert_info = self.validate_certificate(cert_path) + if not cert_info or not cert_info.is_valid: + self.logger.error("Invalid or corrupt certificate") + return False + + cert_dest = self.certs_path / f"{name}_cert.pem" + key_dest = self.certs_path / f"{name}_key.pem" + + shutil.copy2(cert_path, cert_dest) + shutil.copy2(key_path, key_dest) + + # Update config + config = self.proxy_manager.config_manager.get_instance() + config.proxy_settings.gradio_endpoint.cert_name = name + config.proxy_settings.api_endpoint.cert_name = name + config.save() + + self.logger.info(f"Certificate '{name}' installed successfully") + return True + + except Exception as e: + self.logger.error(f"Certificate installation error: {e}") + return False + + def remove_certificate(self, name: str) -> bool: + """Remove a certificate""" + try: + cert_path = self.certs_path / f"{name}_cert.pem" + key_path = self.certs_path / f"{name}_key.pem" + + if cert_path.exists(): + cert_path.unlink() + if key_path.exists(): + key_path.unlink() + + self.logger.info(f"Certificate '{name}' removed successfully") + return True + except Exception as e: + self.logger.error(f"Certificate removal error: {e}") + return False + + def get_certificates_status(self) -> Dict[str, Any]: + """Get status of all installed certificates""" + certificates = {} + for cert_file in self.certs_path.glob("*_cert.pem"): + name = cert_file.name.replace("_cert.pem", "") + cert_info = self.validate_certificate(cert_file) + if cert_info: + certificates[name] = { + "subject": cert_info.subject, + "issuer": cert_info.issuer, + "valid_from": cert_info.valid_from.strftime('%Y-%m-%d %H:%M:%S'), + "valid_until": cert_info.valid_until.strftime('%Y-%m-%d %H:%M:%S'), + "is_valid": cert_info.is_valid, + "days_until_expiry": (cert_info.valid_until - datetime.now()).days + } + return certificates + + def check_expiring_certificates(self, days_warning: int = 30) -> Dict[str, int]: + """Check for certificates nearing expiration""" + expiring = {} + for cert_file in self.certs_path.glob("*_cert.pem"): + name = cert_file.name.replace("_cert.pem", "") + cert_info = self.validate_certificate(cert_file) + if cert_info: + days_left = (cert_info.valid_until - datetime.now()).days + if days_left < days_warning: + expiring[name] = days_left + return expiring \ No newline at end of file diff --git a/system/proxy_module/health_monitor.py b/system/proxy_module/health_monitor.py new file mode 100644 index 0000000..6dc7784 --- /dev/null +++ b/system/proxy_module/health_monitor.py @@ -0,0 +1,127 @@ +import logging +import threading +import time +from typing import Dict, Any +import requests +from dataclasses import dataclass, field +from datetime import datetime + +@dataclass +class HealthMetrics: + last_check: float = 0 + uptime_start: float = 0 + gradio_endpoint_up: bool = False + api_endpoint_up: bool = False + consecutive_failures: int = 0 + total_restarts: int = 0 + response_times: Dict[str, float] = field(default_factory=dict) + +class HealthMonitor: + def __init__(self, proxy_manager): + self.proxy_manager = proxy_manager + self.logger = logging.getLogger("ProxyHealthMonitor") + self.metrics = HealthMetrics() + self.monitoring = False + self.monitor_thread = None + self.check_interval = 30 # seconds + + def start_monitoring(self): + """Start health monitoring thread""" + if self.monitoring: + return + + self.monitoring = True + self.metrics.uptime_start = time.time() + self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True) # Make thread daemon + self.monitor_thread.start() + self.logger.info("Health monitoring started") + + def stop_monitoring(self): + """Stop health monitoring thread""" + if not self.monitoring: + return + + self.monitoring = False + if self.monitor_thread and self.monitor_thread.is_alive(): + try: + # Give the thread a chance to finish naturally + self.monitor_thread.join(timeout=5) + except Exception as e: + self.logger.warning(f"Could not gracefully stop monitoring thread: {e}") + finally: + self.monitor_thread = None + + self.logger.info("Health monitoring stopped") + + def _monitor_loop(self): + """Main monitoring loop""" + while self.monitoring: + self._check_health() + time.sleep(self.check_interval) + + def _check_health(self): + """Check health of proxy endpoints""" + config = self.proxy_manager.config_manager.get_instance() + proxy_settings = config.proxy_settings + self.metrics.last_check = time.time() + + try: + # Only check endpoints that are enabled and running + if proxy_settings.gradio_endpoint.enabled == "Enabled": + try: + start_time = time.time() + response = requests.get( + f"http://127.0.0.1:{config.gradio_port_number}/", # Use internal port + timeout=5 + ) + self.metrics.response_times['gradio'] = time.time() - start_time + self.metrics.gradio_endpoint_up = response.status_code == 200 + except Exception as e: + self.logger.debug(f"Gradio endpoint check failed: {e}") + self.metrics.gradio_endpoint_up = False + + if proxy_settings.api_endpoint.enabled == "Enabled": + try: + start_time = time.time() + response = requests.get( + f"http://127.0.0.1:{config.api_def.api_port_number}/", # Use internal port + timeout=5 + ) + self.metrics.response_times['api'] = time.time() - start_time + self.metrics.api_endpoint_up = response.status_code == 200 + except Exception as e: + self.logger.debug(f"API endpoint check failed: {e}") + self.metrics.api_endpoint_up = False + + if not self.metrics.gradio_endpoint_up and not self.metrics.api_endpoint_up: + self.metrics.consecutive_failures += 1 + else: + self.metrics.consecutive_failures = 0 + + except Exception as e: + self.logger.error(f"Health check failed: {str(e)}") + self.metrics.consecutive_failures += 1 + + # Only attempt restart if we have multiple complete failures + if self.metrics.consecutive_failures >= 3: + self.logger.warning("Too many consecutive failures, attempting restart") + try: + if self.proxy_manager.stop_proxy() and self.proxy_manager.start_proxy(): + self.metrics.total_restarts += 1 + self.metrics.consecutive_failures = 0 + except Exception as e: + self.logger.error(f"Failed to restart proxy: {e}") + + def get_health_status(self) -> Dict[str, Any]: + """Get current health metrics""" + uptime = time.time() - self.metrics.uptime_start if self.metrics.uptime_start else 0 + return { + "status": "healthy" if not self.metrics.consecutive_failures else "degraded", + "uptime": f"{int(uptime/3600)}h {int((uptime%3600)/60)}m", + "last_check": datetime.fromtimestamp(self.metrics.last_check).strftime('%Y-%m-%d %H:%M:%S') + if self.metrics.last_check else "Never", + "gradio_endpoint": "up" if self.metrics.gradio_endpoint_up else "down", + "api_endpoint": "up" if self.metrics.api_endpoint_up else "down", + "response_times": self.metrics.response_times, + "total_restarts": self.metrics.total_restarts + } \ No newline at end of file diff --git a/system/proxy_module/interface.py b/system/proxy_module/interface.py new file mode 100644 index 0000000..777ca3c --- /dev/null +++ b/system/proxy_module/interface.py @@ -0,0 +1,174 @@ +import gradio as gr +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .proxy_manager import ProxyManager + +def create_proxy_interface(proxy_manager: 'ProxyManager'): + """Create the proxy management interface""" + config = proxy_manager.config_manager.get_instance() + + def validate_ports(gradio_port, api_port): + """Validate port configurations""" + if gradio_port == api_port and gradio_port != 0 and api_port != 0: + return False, "Gradio and API ports cannot be the same!" + if not (0 <= gradio_port <= 65535) or not (0 <= api_port <= 65535): + return False, "Ports must be between 0 and 65535" + return True, "Valid port configuration" + + def handle_save_config(proxy_enabled, gradio_enabled, api_enabled, + gradio_port, api_port, startup, external_ip): + # Validate ports first + valid, message = validate_ports(gradio_port, api_port) + if not valid: + return message + + config.proxy_settings.proxy_enabled = proxy_enabled + config.proxy_settings.start_on_startup = startup + + # Update Gradio endpoint + config.proxy_settings.gradio_endpoint.enabled = gradio_enabled + config.proxy_settings.gradio_endpoint.external_port = int(gradio_port) if gradio_enabled else 0 + config.proxy_settings.gradio_endpoint.external_ip = external_ip + + # Update API endpoint + config.proxy_settings.api_endpoint.enabled = api_enabled + config.proxy_settings.api_endpoint.external_port = int(api_port) if api_enabled else 0 + config.proxy_settings.api_endpoint.external_ip = external_ip + + config.save() + return "Configuration saved successfully" + + def update_control_states(proxy_enabled): + """Update enabled/disabled state of controls based on master switch""" + return { + start_button: gr.update(interactive=proxy_enabled), + stop_button: gr.update(interactive=proxy_enabled), + gradio_enabled: gr.update(interactive=proxy_enabled), + api_enabled: gr.update(interactive=proxy_enabled), + start_on_startup: gr.update(interactive=proxy_enabled) + } + + def delete_all_certificates(): + """Delete all certificates and reset config""" + try: + # Delete certificate files + for cert_file in proxy_manager.certs_path.glob("*.pem"): + cert_file.unlink() + + # Reset config + config = proxy_manager.config_manager.get_instance() + config.proxy_settings.gradio_endpoint.cert_name = "" + config.proxy_settings.api_endpoint.cert_name = "" + config.save() + + return "All certificates deleted and config reset" + except Exception as e: + return f"Error deleting certificates: {str(e)}" + + with gr.Tab("Proxy Settings"): + with gr.Column(): + with gr.Row(): + with gr.Group(): + proxy_enabled = gr.Checkbox( + label="Enable Proxy System", + value=config.proxy_settings.proxy_enabled, + info="Master switch/safety lockout for proxy functionality" + ) + start_on_startup = gr.Checkbox( + label="Start Proxy Automatically on Startup", + value=config.proxy_settings.start_on_startup, + info="Automatically start proxy when app launches" + ) + external_ip = gr.Textbox( + label="External IP", + value=config.proxy_settings.gradio_endpoint.external_ip, + info="Use 0.0.0.0 to bind to all interfaces" + ) + with gr.Group(): + gr.Markdown("API interface Proxy Management") + api_enabled = gr.Dropdown( + choices=["Enabled", "Disabled"], + value="Enabled" if config.proxy_settings.api_endpoint.enabled == "Enabled" else "Disabled", + label="API (TTS Generation) Proxy", + info="Enable or Disable Proxying the API (TTS Generation)" + ) + api_port = gr.Number( + value=config.proxy_settings.api_endpoint.external_port, + label="API Port", + info="External port you want to use for the API (TTS Generation)" + ) + with gr.Group(): + gr.Markdown("Gradio interface Proxy Management") + gradio_enabled = gr.Dropdown( + choices=["Enabled", "Disabled"], + value="Enabled" if config.proxy_settings.gradio_endpoint.enabled == "Enabled" else "Disabled", + label="Gradio Interface Proxy", + info="Enable or Disable Proxying the Gradio interface" + ) + gradio_port = gr.Number( + value=config.proxy_settings.gradio_endpoint.external_port, + label="Gradio Interface Port", + info="External port you want to use for the Gradio interface" + ) + + # Status and control section + with gr.Row(): + with gr.Column(scale=2): + status_output = gr.Textbox( + label="Status", + interactive=False, + lines=4 + ) + with gr.Column(scale=1): + save_button = gr.Button("Save Configuration", variant="primary") + with gr.Row(): + start_button = gr.Button("Start Service") + stop_button = gr.Button("Stop Service") + status_button = gr.Button("Check Status") + + # Certificate section + with gr.Row(): + gr.Markdown("### Certificate Management") + with gr.Row(): + with gr.Column(scale=1): + cert_file = gr.File(label="Certificate (.crt/.pem)") + with gr.Column(scale=1): + key_file = gr.File(label="Key (.key)") + with gr.Column(scale=1): + cert_name = gr.Textbox(label="Certificate Name", placeholder="e.g., gradio_cert") + upload_button = gr.Button("Upload Certificate", variant="primary") + delete_button = gr.Button("Delete Certificates", variant="secondary") + + # Connect everything + save_button.click( + fn=handle_save_config, + inputs=[ + proxy_enabled, + gradio_enabled, + api_enabled, + gradio_port, + api_port, + start_on_startup, + external_ip + ], + outputs=status_output + ) + proxy_enabled.change( + fn=update_control_states, + inputs=[proxy_enabled], + outputs=[start_button, stop_button, gradio_enabled, api_enabled, start_on_startup] + ) + start_button.click(fn=proxy_manager.start_proxy, outputs=status_output) + stop_button.click(fn=proxy_manager.stop_proxy, outputs=status_output) + status_button.click(fn=proxy_manager.get_status, outputs=status_output) + upload_button.click( + fn=proxy_manager.handle_cert_upload, + inputs=[cert_file, key_file, cert_name], + outputs=status_output + ) + delete_button.click(fn=delete_all_certificates, outputs=status_output) + + return proxy_enabled, start_on_startup, gradio_port, api_port \ No newline at end of file diff --git a/system/proxy_module/metrics.py b/system/proxy_module/metrics.py new file mode 100644 index 0000000..a3541a4 --- /dev/null +++ b/system/proxy_module/metrics.py @@ -0,0 +1,94 @@ +import logging +import threading +import time +from typing import Dict, Any +from collections import defaultdict +from dataclasses import dataclass, field +from datetime import datetime, timedelta + +@dataclass +class ProxyMetrics: + start_time: float = field(default_factory=time.time) + total_requests: int = 0 + active_connections: int = 0 + bytes_transferred: int = 0 + request_times: Dict[str, float] = field(default_factory=lambda: defaultdict(float)) + requests_per_minute: Dict[str, int] = field(default_factory=lambda: defaultdict(int)) + status_codes: Dict[str, int] = field(default_factory=lambda: defaultdict(int)) + +class MetricsCollector: + def __init__(self, proxy_manager): + self.proxy_manager = proxy_manager + self.logger = logging.getLogger("ProxyMetrics") + self.metrics = ProxyMetrics() + self.collecting = False + self.collector_thread = None + self.collect_interval = 60 # seconds + self._lock = threading.Lock() + + def start_collecting(self): + """Start metrics collection""" + if self.collecting: + return + + self.collecting = True + self.metrics = ProxyMetrics() # Reset metrics + self.collector_thread = threading.Thread(target=self._collect_loop) + self.collector_thread.daemon = True + self.collector_thread.start() + self.logger.info("Metrics collection started") + + def stop_collecting(self): + """Stop metrics collection""" + self.collecting = False + if self.collector_thread: + self.collector_thread.join() + self.logger.info("Metrics collection stopped") + + def _collect_loop(self): + """Main collection loop""" + while self.collecting: + self._update_metrics() + time.sleep(self.collect_interval) + + def _update_metrics(self): + """Update periodic metrics""" + with self._lock: + # Reset per-minute counters + self.metrics.requests_per_minute.clear() + + def record_request(self, endpoint: str, status_code: int, response_time: float, bytes_sent: int): + """Record metrics for a single request""" + with self._lock: + self.metrics.total_requests += 1 + self.metrics.bytes_transferred += bytes_sent + self.metrics.request_times[endpoint] = response_time + self.metrics.status_codes[str(status_code)] += 1 + self.metrics.requests_per_minute[endpoint] += 1 + + def get_current_metrics(self) -> Dict[str, Any]: + """Get current metrics""" + with self._lock: + uptime = time.time() - self.metrics.start_time + return { + "uptime": str(timedelta(seconds=int(uptime))), + "total_requests": self.metrics.total_requests, + "active_connections": self.metrics.active_connections, + "bytes_transferred": f"{self.metrics.bytes_transferred / 1024 / 1024:.2f} MB", + "average_response_times": { + endpoint: f"{time:.2f}ms" + for endpoint, time in self.metrics.request_times.items() + }, + "requests_per_minute": dict(self.metrics.requests_per_minute), + "status_codes": dict(self.metrics.status_codes) + } + + def increment_connections(self): + """Increment active connection count""" + with self._lock: + self.metrics.active_connections += 1 + + def decrement_connections(self): + """Decrement active connection count""" + with self._lock: + self.metrics.active_connections = max(0, self.metrics.active_connections - 1) \ No newline at end of file diff --git a/system/proxy_module/proxy_manager.py b/system/proxy_module/proxy_manager.py new file mode 100644 index 0000000..f63a405 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/system/proxy_module/security.py b/system/proxy_module/security.py new file mode 100644 index 0000000..3116184 --- /dev/null +++ b/system/proxy_module/security.py @@ -0,0 +1,61 @@ +import logging +from typing import Dict, Any, Set +import time +from collections import defaultdict +from dataclasses import dataclass +from threading import Lock + +@dataclass +class SecurityMetrics: + total_requests: int = 0 + blocked_requests: int = 0 + last_incident: float = 0 + blacklisted_ips: Set[str] = None + + def __post_init__(self): + if self.blacklisted_ips is None: + self.blacklisted_ips = set() + +class SecurityManager: + def __init__(self, proxy_manager): + self.proxy_manager = proxy_manager + self.logger = logging.getLogger("ProxySecurityManager") + self.metrics = SecurityMetrics() + self.rate_limits = defaultdict(lambda: {"count": 0, "reset_time": 0}) + self.rate_limit_lock = Lock() + + def check_rate_limit(self, ip: str, limit: int = 100, window: int = 60) -> bool: + """Check if IP has exceeded rate limit""" + with self.rate_limit_lock: + current_time = time.time() + if self.rate_limits[ip]["reset_time"] < current_time: + self.rate_limits[ip] = {"count": 0, "reset_time": current_time + window} + + self.rate_limits[ip]["count"] += 1 + return self.rate_limits[ip]["count"] <= limit + + def is_blacklisted(self, ip: str) -> bool: + """Check if IP is blacklisted""" + return ip in self.metrics.blacklisted_ips + + def add_to_blacklist(self, ip: str, reason: str): + """Add IP to blacklist""" + self.metrics.blacklisted_ips.add(ip) + self.logger.warning(f"IP {ip} blacklisted: {reason}") + self.metrics.last_incident = time.time() + + def remove_from_blacklist(self, ip: str): + """Remove IP from blacklist""" + self.metrics.blacklisted_ips.discard(ip) + self.logger.info(f"IP {ip} removed from blacklist") + + def get_security_status(self) -> Dict[str, Any]: + """Get current security metrics""" + return { + "total_requests": self.metrics.total_requests, + "blocked_requests": self.metrics.blocked_requests, + "blacklisted_ips": len(self.metrics.blacklisted_ips), + "last_incident": time.strftime('%Y-%m-%d %H:%M:%S', + time.localtime(self.metrics.last_incident)) + if self.metrics.last_incident else "None" + } \ No newline at end of file diff --git a/system/proxy_module/twisted_server.py b/system/proxy_module/twisted_server.py new file mode 100644 index 0000000..bbb4987 --- /dev/null +++ b/system/proxy_module/twisted_server.py @@ -0,0 +1,218 @@ +import subprocess +import sys +from pathlib import Path +import logging + +# Auto-install requirements if missing +def install_requirements(): + """Install missing requirements only if they're not already in sys.modules""" + required_packages = [ + "twisted", + "cryptography", + "service_identity", + "requests", + "psutil" + ] + + for package in required_packages: + if package not in sys.modules: + try: + __import__(package) + except ImportError: + print(f"Package '{package}' is not installed. Installing now...") + subprocess.check_call([sys.executable, "-m", "pip", "install", package]) + +install_requirements() + +from twisted.internet import reactor, ssl +from twisted.web import proxy, server +import argparse +import inspect + +# Handle config import for both package and standalone usage +try: + from ...config import AlltalkConfig # Package import +except ImportError: + # Get the path to the root directory (3 levels up from this file) + root_dir = Path(__file__).parent.parent.parent + sys.path.append(str(root_dir)) + from config import AlltalkConfig # Standalone import + +def initialize_configs(): + """Initialize all configuration instances""" + config_initalize = AlltalkConfig.get_instance() + return config_initalize + +# Load in configs +central_config = initialize_configs() + +########################## +# Central print function # +########################## +# ANSI color codes +BLUE = "\033[94m" +# MAGENTA = "\033[95m" +YELLOW = "\033[93m" +RED = "\033[91m" +GREEN = "\033[92m" +RESET = "\033[0m" + +def print_message(message, message_type="standard", component="PRX"): + """Centralized print function for AllTalk messages + Args: + message (str): The message to print + message_type (str): Type of message (standard/warning/error/debug_*/debug) + component (str): Component identifier (TTS/ENG/GEN/API/etc.) + """ + prefix = f"[{central_config.branding}{component}] " + + if message_type.startswith("debug_"): + debug_flag = getattr(central_config.debugging, message_type, False) + if not debug_flag: + return + + if message_type == "debug_func" and "Function entry:" in message: + message_parts = message.split("Function entry:", 1) + print( + f"{prefix}{BLUE}Debug{RESET} {YELLOW}{message_type}{RESET} Function entry:{GREEN}{message_parts[1]}{RESET} twisted_server.py" + ) + else: + print(f"{prefix}{BLUE}Debug{RESET} {YELLOW}{message_type}{RESET} {message}") + + elif message_type == "debug": + print(f"{prefix}{BLUE}Debug{RESET} {message}") + + elif message_type == "warning": + print(f"{prefix}{YELLOW}Warning{RESET} {message}") + + elif message_type == "error": + print(f"{prefix}{RED}Error{RESET} {message}") + + else: + print(f"{prefix}{message}") + +def debug_func_entry(): + """Print debug message for function entry if debug_func is enabled""" + if central_config.debugging.debug_func: + current_func = inspect.currentframe().f_back.f_code.co_name + print_message(f"Function entry: {current_func}", "debug_func") + +class ProxyToInternalPort(proxy.ReverseProxyResource): + def __init__(self, target_host, target_port): + super().__init__(target_host, target_port, b"") + self.logger = logging.getLogger("ProxyManager") + + def proxyClientConnectionFailed(self, connector, reason): + print_message(f"Connection failed: {reason}", message_type="debug_proxy") + return proxy.ReverseProxyResource.proxyClientConnectionFailed(self, connector, reason) + + def proxyClientConnectionMade(self, connector): + print_message(f"Connection established to {self.host}:{self.port}", message_type="debug_proxy") + return proxy.ReverseProxyResource.proxyClientConnectionMade(self, connector) + + def getChild(self, path, request): + try: + decoded_path = path.decode('utf-8', 'ignore') + client_addr = request.getClientAddress() + ip = client_addr.host if hasattr(client_addr, 'host') else 'unknown' + method = request.method.decode('utf-8', 'ignore') + + print_message(f"{GREEN}Request: {RED}{method} {decoded_path}{GREEN} from {RESET}{ip} {GREEN}Forwarding to:{RESET} {self.host}:{self.port} {RED}{decoded_path}{RESET}", message_type="debug_proxy") + #print_message(f"Forwarding to: {self.host}:{self.port}{decoded_path}", message_type="debug_proxy") + except Exception as e: + print_message(f"Error in getChild: {str(e)}", message_type="error") + + child = proxy.ReverseProxyResource(self.host, self.port, self.path + b'/' + path) + return child + +def check_certificates(cert_file_base): + """Check if both key and cert files exist for the given certificate base name.""" + debug_func_entry() + if cert_file_base != "0": + cert_path = Path("certificates") # Add base path + key_file = cert_path / f"{cert_file_base}_key.pem" + cert_file = cert_path / f"{cert_file_base}_cert.pem" + + if not key_file.exists() or not cert_file.exists(): + return False + return True + return False + +def run_proxy(external_ip_a, internal_ip_a, external_ip_b, internal_ip_b, cert_file_a, cert_file_b): + debug_func_entry() + proxies_started = False + api_port = central_config.api_def.api_port_number + gradio_port = central_config.gradio_port_number + api_port_external = central_config.proxy_settings.api_endpoint.external_port + gradio_port_external = central_config.proxy_settings.gradio_endpoint.external_port + + print_message("Starting proxy setup with the following parameters:", message_type="debug_proxy") + print_message(f" |- External IP API : {'[DISABLED]' if external_ip_a == '0' else f'{external_ip_a}:{api_port_external}'} -> {internal_ip_a}:{api_port}", message_type="debug_proxy") + print_message(f" |- External IP Gradio : {'[DISABLED]' if external_ip_b == '0' else f'{external_ip_b}:{gradio_port_external}'} -> {internal_ip_b}:{gradio_port}", message_type="debug_proxy") + print_message(f" |- Certificate File A : {cert_file_a} {'[No Certificate So Disabled]' if cert_file_a == '0' else ''}", message_type="debug_proxy") + print_message(f" |- Certificate File B : {cert_file_b} {'[No Certificate So Disabled]' if cert_file_b == '0' else ''}", message_type="debug_proxy") + + # Gradio Proxy + if external_ip_b != "0" and central_config.proxy_settings.gradio_endpoint.enabled == "Enabled": + if cert_file_b != "0" and check_certificates(cert_file_b): + sslContextB = ssl.DefaultOpenSSLContextFactory( + f"certificates/{cert_file_b}_key.pem", + f"certificates/{cert_file_b}_cert.pem" + ) + proxy_gradio = ProxyToInternalPort(internal_ip_b, gradio_port) + proxy_site_gradio = server.Site(proxy_gradio) + reactor.listenSSL(gradio_port_external, proxy_site_gradio, sslContextB, interface=external_ip_b) + print_message(f"Gradio SSL Proxy : {GREEN}From{RESET} {external_ip_b}:{gradio_port_external} {GREEN}->{RESET} {internal_ip_b}:{gradio_port} {GREEN}is now active.{RESET}") + else: + proxy_gradio = ProxyToInternalPort(internal_ip_b, gradio_port) + proxy_site_gradio = server.Site(proxy_gradio) + reactor.listenTCP(gradio_port_external, proxy_site_gradio, interface=external_ip_b) + print_message(f"Gradio HTTP Proxy : {GREEN}From{RESET} {external_ip_b}:{gradio_port_external} {GREEN}->{RESET} {internal_ip_b}:{gradio_port} {GREEN}is now active.{RESET}") + proxies_started = True + else: + print_message(f"Gradio Proxy is {RED}DISABLED{RESET}") + + # API Proxy + if external_ip_a != "0" and central_config.proxy_settings.api_endpoint.enabled == "Enabled": + if cert_file_a != "0" and check_certificates(cert_file_a): + sslContextA = ssl.DefaultOpenSSLContextFactory( + f"certificates/{cert_file_a}_key.pem", + f"certificates/{cert_file_a}_cert.pem" + ) + proxy_api = ProxyToInternalPort(internal_ip_a, api_port) + proxy_site_api = server.Site(proxy_api) + reactor.listenSSL(api_port_external, proxy_site_api, sslContextA, interface=external_ip_a) + print_message(f"API SSL Proxy : {GREEN}From{RESET} {external_ip_a}:{api_port_external} {GREEN}->{RESET} {internal_ip_a}:{api_port} {GREEN}is now active.{RESET}") + else: + proxy_api = ProxyToInternalPort(internal_ip_a, api_port) + proxy_site_api = server.Site(proxy_api) + reactor.listenTCP(api_port_external, proxy_site_api, interface=external_ip_a) + print_message(f"API HTTP Proxy : {GREEN}From{RESET} {external_ip_a}:{api_port_external} {GREEN}->{RESET} {internal_ip_a}:{api_port} {GREEN}is now active.{RESET}") + proxies_started = True + else: + print_message(f"API Proxy is {RED}DISABLED{RESET}") + + if not proxies_started: + print_message(f"No proxies were started. {RED}Both{RESET} API and Gradio proxies are {RED}DISABLED{RESET}.", message_type="warning") + return + + reactor.run() + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--external_ip_a", required=True, help="External IP Address A ('0' to disable)") + parser.add_argument("--internal_ip_a", required=True, help="Internal IP Address A") + parser.add_argument("--external_ip_b", required=True, help="External IP Address B ('0' to disable)") + parser.add_argument("--internal_ip_b", required=True, help="Internal IP Address B") + parser.add_argument("--cert_file_a", required=True, help="Certificate file base name for External IP A (without '_key' or '_cert', or '0' for HTTP only)") + parser.add_argument("--cert_file_b", required=True, help="Certificate file base name for External IP B (without '_key' or '_cert', or '0' for HTTP only)") + args = parser.parse_args() + + run_proxy( + external_ip_a=args.external_ip_a, + internal_ip_a=args.internal_ip_a, + external_ip_b=args.external_ip_b, + internal_ip_b=args.internal_ip_b, + cert_file_a=args.cert_file_a, + cert_file_b=args.cert_file_b + )