NexusOS/synapse/proc_util.py

"""Process inspection and termination that works with or without psutil.

psutil became an optional extra when the wheel landed, so the base install
(notably `pip install nexusos-ai` on Windows) has to manage PIDs with the
stdlib alone. Callers should keep using psutil when it is importable - it is
faster and more precise - and fall back here when it is not.

Windows has no signals: os.kill(pid, sig) special-cases CTRL_C_EVENT and
CTRL_BREAK_EVENT, treats sig 0 as an existence check, and calls
TerminateProcess(handle, sig) for *everything else*. So os.kill(pid, 15) is not
a polite request there - it is an immediate, unblockable kill with exit code
15, and there is no equivalent of SIGTERM. Everything below goes through the
Win32 API via ctypes so the intent is explicit at each call site rather than
resting on which signal numbers happen to be special.
"""
from __future__ import annotations

import os
import signal
import subprocess
from pathlib import Path

WINDOWS = os.name == "nt"

# Win32 constants (winnt.h / processthreadsapi.h)
_SYNCHRONIZE = 0x00100000
_PROCESS_TERMINATE = 0x0001
_WAIT_TIMEOUT = 0x00000102

# Keep console windows from flashing on every helper subprocess.
_NO_WINDOW = subprocess.CREATE_NO_WINDOW if WINDOWS else 0


def _kernel32():
    import ctypes

    return ctypes.WinDLL("kernel32", use_last_error=True)


def _run(argv: list[str], timeout: float = 10.0) -> str:
    """Run a helper command, returning stdout ('' on any failure)."""
    try:
        done = subprocess.run(
            argv, capture_output=True, text=True, timeout=timeout,
            creationflags=_NO_WINDOW,
        )
    except (OSError, subprocess.SubprocessError):
        return ""
    return done.stdout or ""


def pid_alive(pid: int | None) -> bool:
    """True if the PID names a live process.

    On Windows this opens a handle and polls it; a signalled handle means the
    process has exited. That is equivalent to os.kill(pid, 0) - CPython
    special-cases signal 0 into an existence check there - but it also reports
    True for a process owned by another user, where os.kill raises.
    """
    if pid is None:
        return False
    try:
        pid = int(pid)
    except (TypeError, ValueError):
        return False
    if pid <= 0:
        return False
    if WINDOWS:
        import ctypes

        k = _kernel32()
        handle = k.OpenProcess(_SYNCHRONIZE, False, pid)
        if not handle:
            return False
        try:
            return k.WaitForSingleObject(ctypes.c_void_p(handle), 0) == _WAIT_TIMEOUT
        finally:
            k.CloseHandle(ctypes.c_void_p(handle))
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True          # exists, owned by someone else
    except (OSError, ValueError):
        return False


def terminate_pid(pid: int | None, force: bool = False) -> bool:
    """Ask a process to exit. Returns True if the request was delivered."""
    if not pid_alive(pid):
        return False
    pid = int(pid)
    if WINDOWS:
        # No graceful path without a shared console; TerminateProcess is what
        # psutil.terminate() resolves to on Windows anyway.
        import ctypes

        k = _kernel32()
        handle = k.OpenProcess(_PROCESS_TERMINATE, False, pid)
        if not handle:
            return False
        try:
            return bool(k.TerminateProcess(ctypes.c_void_p(handle), 1))
        finally:
            k.CloseHandle(ctypes.c_void_p(handle))
    try:
        os.kill(pid, signal.SIGKILL if force else signal.SIGTERM)
        return True
    except OSError:
        return False


def pid_cmdline(pid: int | None) -> str:
    """Full command line for a PID, or '' when it cannot be determined."""
    if pid is None:
        return ""
    try:
        pid = int(pid)
    except (TypeError, ValueError):
        return ""
    if WINDOWS:
        for entry_pid, cmd in iter_processes():
            if entry_pid == pid:
                return cmd
        return ""
    proc = Path(f"/proc/{pid}/cmdline")
    try:
        return proc.read_bytes().replace(b"\0", b" ").decode(errors="replace").strip()
    except OSError:
        pass
    # macOS and other POSIX hosts without /proc.
    out = _run(["ps", "-o", "command=", "-p", str(pid)])
    return out.strip()


def iter_processes() -> list[tuple[int, str]]:
    """(pid, command_line) for every visible process.

    Windows needs CIM for command lines - the ctypes snapshot APIs only expose
    image names, which is not enough to tell `uvicorn synapse.main` apart from
    any other python.exe. This is slow, so it is strictly the no-psutil path.
    """
    if WINDOWS:
        out = _run([
            "powershell", "-NoProfile", "-NonInteractive", "-Command",
            "Get-CimInstance Win32_Process | "
            "ForEach-Object { \"$($_.ProcessId)`t$($_.CommandLine)\" }",
        ], timeout=30.0)
        entries: list[tuple[int, str]] = []
        for line in out.splitlines():
            head, _, cmd = line.partition("\t")
            if head.strip().isdigit():
                entries.append((int(head), cmd.strip()))
        return entries

    entries = []
    proc_root = Path("/proc")
    if proc_root.is_dir():
        for entry in proc_root.iterdir():
            if not entry.name.isdigit():
                continue
            try:
                cmd = (entry / "cmdline").read_bytes()
            except OSError:
                continue
            entries.append((int(entry.name), cmd.replace(b"\0", b" ").decode(errors="replace").strip()))
        return entries

    for line in _run(["ps", "-A", "-o", "pid=,command="]).splitlines():
        head, _, cmd = line.strip().partition(" ")
        if head.isdigit():
            entries.append((int(head), cmd.strip()))
    return entries


def pids_listening_on(port: int) -> list[int]:
    """PIDs holding a listening TCP socket on `port`."""
    pids: list[int] = []
    if WINDOWS:
        for line in _run(["netstat", "-ano", "-p", "TCP"]).splitlines():
            parts = line.split()
            # Proto  Local            Foreign  State      PID
            if len(parts) < 5 or parts[3] != "LISTENING":
                continue
            local = parts[1]
            if local.rsplit(":", 1)[-1] == str(port) and parts[4].isdigit():
                pids.append(int(parts[4]))
        return sorted(set(pids))

    # -t TCP, -l listening, -n numeric, -P no port names.
    for line in _run(["lsof", "-nP", "-tiTCP:%d" % port, "-sTCP:LISTEN"]).splitlines():
        if line.strip().isdigit():
            pids.append(int(line.strip()))
    if pids:
        return sorted(set(pids))

    for line in _run(["ss", "-lptnH", "sport = :%d" % port]).splitlines():
        # ... users:(("uvicorn",pid=1234,fd=3))
        marker = "pid="
        start = line.find(marker)
        while start != -1:
            digits = ""
            for ch in line[start + len(marker):]:
                if not ch.isdigit():
                    break
                digits += ch
            if digits:
                pids.append(int(digits))
            start = line.find(marker, start + 1)
    return sorted(set(pids))