#!/usr/bin/env python3
"""
Per-connection noVNC guard (launchd socket activation / inetd mode).

Before each WebSocket connection, kill ALL existing VNC client handlers on
port 5900. Apple Screen Sharing allows only one session.
"""

import os
import signal
import subprocess
import sys
import time

VNC_PORT = 5900
WEBSOCKIFY = "/Users/agent/Library/Python/3.9/bin/websockify"
WEB_DIR = "/Users/agent/noVNC"
LOG_FILE = "/tmp/websockify-guard.log"


def log(msg):
    with open(LOG_FILE, "a") as f:
        f.write(f"{time.strftime('%H:%M:%S')} {msg}\n")


def kill_all_vnc_clients():
    """Kill every ESTABLISHED client to port 5900 (never screensharingd)."""
    killed = []
    try:
        result = subprocess.run(
            ["lsof", "-i", f":{VNC_PORT}", "-n", "-P"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        for line in result.stdout.strip().split("\n")[1:]:
            if "ESTABLISHED" not in line or "->" not in line:
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            pid = int(parts[1])
            try:
                ps = subprocess.run(
                    ["ps", "-p", str(pid), "-o", "command="],
                    capture_output=True,
                    text=True,
                    timeout=3,
                )
                if "screensharingd" in ps.stdout:
                    continue
            except Exception:
                pass
            try:
                os.kill(pid, signal.SIGTERM)
                killed.append(pid)
            except (ProcessLookupError, PermissionError):
                pass
        if killed:
            log(f"Pre-connect cleanup: killed PIDs {killed}")
            time.sleep(0.3)
    except Exception as e:
        log(f"cleanup error: {e}")


if __name__ == "__main__":
    kill_all_vnc_clients()
    os.execvp(
        WEBSOCKIFY,
        [
            "websockify",
            "--web",
            WEB_DIR,
            "--heartbeat",
            "15",
            "--idle-timeout",
            "0",
            "--inetd",
            f"localhost:{VNC_PORT}",
        ],
    )