235 lines
9.4 KiB
Python
235 lines
9.4 KiB
Python
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") |