THeGoVerNor212/litellm

1#!/usr/bin/env python32"""3Simple Python-based reverse proxy with multi-process management4No supervisord needed - runs nginx and sshd as subprocesses5"""6import os7import signal8import subprocess9import sys10import time11from pathlib import Path12 13# Track running processes14processes = []15 16 17def signal_handler(sig, frame):18    """Handle shutdown signals gracefully"""19    print(f"\nReceived signal {sig}, shutting down gracefully...")20    for proc in processes:21        try:22            proc.terminate()23            proc.wait(timeout=5)24        except subprocess.TimeoutExpired:25            proc.kill()26        except Exception as e:27            print(f"Error stopping process: {e}")28    sys.exit(0)29 30 31def setup_ssh():32    """Generate SSH host keys if they don't exist"""33    key_dir = Path("/home/user/.ssh/hostkeys")34    key_dir.mkdir(parents=True, exist_ok=True)35 36    rsa_key = key_dir / "ssh_host_rsa_key"37    ed_key = key_dir / "ssh_host_ed25519_key"38 39    if not rsa_key.exists() or not ed_key.exists():40        print("Generating SSH host keys...")41        subprocess.run(42            ["ssh-keygen", "-t", "rsa", "-b", "4096", "-f", str(rsa_key), "-N", ""],43            check=True,44        )45        subprocess.run(46            ["ssh-keygen", "-t", "ed25519", "-f", str(ed_key), "-N", ""], check=True47        )48    else:49        print("SSH host keys already exist")50 51    # Add SSH public key from environment if provided52    ssh_public_key = os.getenv("SSH_PUBLIC_KEY")53    if ssh_public_key:54        print("Adding SSH public key from environment...")55        ssh_dir = Path("/home/user/.ssh")56        ssh_dir.mkdir(parents=True, exist_ok=True)57 58        authorized_keys = ssh_dir / "authorized_keys"59        with open(authorized_keys, "a") as f:60            f.write(f"{ssh_public_key}\n")61 62        os.chmod(ssh_dir, 0o700)63        os.chmod(authorized_keys, 0o600)64        subprocess.run(["chown", "-R", "user:user", str(ssh_dir)], check=True)65        print("SSH public key added successfully")66 67 68def start_nginx():69    """Start nginx in foreground mode"""70    print("Starting nginx...")71    proc = subprocess.Popen(72        ["/usr/sbin/nginx", "-g", "daemon off;"],73        stdout=subprocess.PIPE,74        stderr=subprocess.STDOUT,75        text=True,76    )77    processes.append(proc)78    return proc79 80 81def start_sshd():82    """Start sshd in non-daemon mode"""83    print("Starting sshd...")84    proc = subprocess.Popen(85        ["/usr/sbin/sshd", "-D"],86        stdout=subprocess.PIPE,87        stderr=subprocess.STDOUT,88        text=True,89    )90    processes.append(proc)91    return proc92 93 94def start_tailscale():95    """Start Tailscale if auth key is provided"""96    auth_key = os.getenv("TAILSCALE_AUTHKEY")97    if not auth_key:98        print("No TAILSCALE_AUTHKEY provided, skipping Tailscale")99        return None, None100 101    print("Starting tailscaled...")102    tailscaled = subprocess.Popen(103        [104            "/usr/sbin/tailscaled",105            "--state=/var/lib/tailscale/tailscaled.state",106            "--socket=/var/run/tailscale/tailscaled.sock",107            "--tun=userspace-networking",108        ],109        stdout=subprocess.PIPE,110        stderr=subprocess.STDOUT,111        text=True,112    )113    processes.append(tailscaled)114 115    # Wait for tailscaled to be ready116    print("Waiting for tailscaled to start...")117    for i in range(30):118        result = subprocess.run(119            ["tailscale", "status"],120            stdout=subprocess.DEVNULL,121            stderr=subprocess.DEVNULL,122        )123        if result.returncode == 0:124            break125        time.sleep(1)126    else:127        print("Warning: tailscaled did not start in time")128        return tailscaled, None129 130    # Connect to Tailscale network131    print("Connecting to Tailscale network...")132    tailscale_up = subprocess.Popen(133        ["tailscale", "up", "--authkey", auth_key, "--accept-routes"],134        stdout=subprocess.PIPE,135        stderr=subprocess.STDOUT,136        text=True,137    )138    tailscale_up.wait()139    print("Tailscale connection initiated")140 141    return tailscaled, tailscale_up142 143 144def monitor_processes():145    """Monitor all processes and restart if they crash"""146    restart_count = {}147    max_restarts = 3148    restart_window = 60  # seconds149 150    while True:151        time.sleep(5)152 153        for i, proc in enumerate(processes[:]):  # Copy list to iterate safely154            if proc.poll() is not None:  # Process has died155                proc_name = {0: "nginx", 1: "sshd", 2: "tailscaled"}.get(156                    i, f"process-{i}"157                )158 159                print(160                    f"WARNING: {proc_name} (PID {proc.pid}) has stopped with code {proc.returncode}"161                )162 163                # Check restart limits164                now = time.time()165                if proc_name not in restart_count:166                    restart_count[proc_name] = []167 168                # Remove old restart timestamps169                restart_count[proc_name] = [170                    ts for ts in restart_count[proc_name] if now - ts < restart_window171                ]172 173                if len(restart_count[proc_name]) >= max_restarts:174                    print(f"ERROR: {proc_name} has restarted too many times, giving up")175                    continue176 177                # Restart the process178                print(f"Restarting {proc_name}...")179                restart_count[proc_name].append(now)180 181                if i == 0:  # nginx182                    new_proc = start_nginx()183                    processes[i] = new_proc184                elif i == 1:  # sshd185                    new_proc = start_sshd()186                    processes[i] = new_proc187                # Don't auto-restart tailscaled as it's more complex188 189 190def main():191    """Main entry point"""192    print("=" * 60)193    print("Starting Multi-Process Service Manager")194    print("=" * 60)195 196    # Register signal handlers197    signal.signal(signal.SIGTERM, signal_handler)198    signal.signal(signal.SIGINT, signal_handler)199 200    # Setup SSH201    setup_ssh()202 203    # Start Tailscale (optional)204    start_tailscale()205 206    # Start core services207    start_nginx()208    start_sshd()209 210    print("\n" + "=" * 60)211    print("All services started successfully!")212    print("=" * 60)213    print(f"Running with {len(processes)} processes")214    print("Press Ctrl+C to stop\n")215 216    # Monitor processes indefinitely217    try:218        monitor_processes()219    except KeyboardInterrupt:220        signal_handler(signal.SIGINT, None)221 222 223if __name__ == "__main__":224    main()225