#!/usr/bin/env python3
# Copyright (c) 2026 Storm Dragon
# Authors: Storm Dragon, Angus Kola
# SPDX-License-Identifier: GPL-3.0-or-later
"""Crash the real quarantined mailbox owner, only inside native_remote_keys.py."""
import ctypes as C
import os
from pathlib import Path
import signal
import subprocess
import tempfile
import time


def key_state():
    x = C.CDLL("libX11.so.6")
    x.XOpenDisplay.argtypes, x.XOpenDisplay.restype = [C.c_char_p], C.c_void_p
    x.XStringToKeysym.argtypes, x.XStringToKeysym.restype = [C.c_char_p], C.c_ulong
    x.XKeysymToKeycode.argtypes, x.XKeysymToKeycode.restype = [C.c_void_p, C.c_ulong], C.c_ubyte
    x.XQueryKeymap.argtypes = [C.c_void_p, C.c_void_p]
    x.XCloseDisplay.argtypes = [C.c_void_p]
    display = x.XOpenDisplay(None)
    if not display:
        raise RuntimeError("Private X key-state connection unavailable")
    try:
        keys = (C.c_ubyte * 32)()
        if not x.XQueryKeymap(display, keys):
            raise RuntimeError("Private X key-state query failed")
        result = []
        for name in (b"Shift_L", b"Control_R", b"a"):
            code = x.XKeysymToKeycode(display, x.XStringToKeysym(name))
            if not code:
                raise RuntimeError("Private test key is unavailable")
            result.append(bool(keys[code // 8] & (1 << (code % 8))))
        return tuple(result)
    finally:
        x.XCloseDisplay(display)


def wait_for(predicate, message, timeout=7):
    deadline = time.monotonic() + timeout
    while not predicate():
        if time.monotonic() >= deadline:
            raise RuntimeError(message)
        time.sleep(0.01)


def children(pid):
    result = set()
    try:
        tasks = list(Path(f"/proc/{pid}/task").iterdir())
    except FileNotFoundError:
        return result
    for task in tasks:
        try:
            result.update(int(value) for value in (task / "children").read_text().split())
        except FileNotFoundError:
            pass  # A thread may exit between the directory and children reads.
    return result


def scenario(mode):
    with tempfile.TemporaryDirectory(prefix="cthulhu-owner-crash-") as temporary:
        ready = Path(temporary) / "ready"
        environment = dict(os.environ, CTHULHU_REMOTE_OWNER_READY=str(ready), TMPDIR=temporary)
        with tempfile.TemporaryFile(mode="w+b") as log:
            owner = subprocess.Popen(
                [environment["CTHULHU_REMOTE_OWNER_TEST"],
                 "-test.run=^TestNativeRemoteKeyboardOwnerProcess$", "-test.timeout=30s", "-test.v"],
                env=environment, stdin=subprocess.PIPE, stdout=log, stderr=log,
                start_new_session=True)
            descendants = set()
            try:
                def is_ready():
                    if owner.poll() is not None:
                        raise RuntimeError("Reader-side owner exited before readiness")
                    return ready.exists()

                wait_for(is_ready, "Reader-side owner did not hold both modifiers", timeout=12)
                if key_state() != (True, True, False):
                    raise RuntimeError("Readiness did not retain the expected native keys")
                descendants = children(owner.pid)
                helper_path = Path(environment["CTHULHU_REMOTE_INPUT_HELPER"]).resolve()
                helpers = [pid for pid in descendants if Path(f"/proc/{pid}/exe").resolve() == helper_path]
                if len(helpers) != 1 or len(descendants) != 2:
                    raise RuntimeError("Expected exactly the source helper and SDK peer")
                helper = helpers[0]
                if os.getpgid(helper) == os.getpgid(owner.pid):
                    raise RuntimeError("Helper shared the crashed reader's process group")
                began = time.monotonic()
                # Crash/stop the reader's entire group, preserving the helper's
                # independent group. No graceful cancellation runs in the owner.
                os.killpg(owner.pid, signal.SIGKILL if mode == "kill" else signal.SIGSTOP)
                wait_for(lambda: key_state() == (False, False, False),
                         "Independent helper did not release the crashed owner's keys")
                elapsed = time.monotonic() - began
                if mode == "stop":
                    # A stopped owner keeps its lifetime descriptor open: only
                    # the helper's independent heartbeat watchdog can retire it.
                    if elapsed < 2:
                        raise RuntimeError("Stopped-owner recovery preceded watchdog expiry")
                    os.killpg(owner.pid, signal.SIGCONT)
                    owner.stdin.write(b"c")
                    owner.stdin.flush()
                    if owner.wait(timeout=10) != 0:
                        raise RuntimeError("Resumed owner did not retain retirement barriers")
                else:
                    if owner.wait(timeout=5) != -signal.SIGKILL:
                        raise RuntimeError("Reader-side owner was not killed as intended")

                    def reaped():
                        for pid in list(descendants):
                            ended, _ = os.waitpid(pid, os.WNOHANG)
                            if ended:
                                descendants.remove(pid)
                        return not descendants

                    wait_for(reaped, "Orphaned helper or SDK peer did not exit")
                if key_state() != (False, False, False):
                    raise RuntimeError("Retired owner restored native key state")
                print(f"PASS mailbox owner SIG{mode.upper()}: keys released in {elapsed:.2f}s", flush=True)
            except BaseException:
                log.seek(0)
                print(log.read().decode(errors="replace"), flush=True)
                raise
            finally:
                # Only disposable fixture processes are killed here. Production
                # retirement never kills a helper to satisfy a timeout.
                if owner.poll() is None:
                    descendants.update(children(owner.pid))
                    os.killpg(owner.pid, signal.SIGKILL)
                    owner.wait(timeout=5)
                owner.stdin.close()
                # The resumed Go owner may already have reaped these PIDs.
                # Signal only descendants currently adopted by this supervisor.
                for pid in descendants & children(os.getpid()):
                    try:
                        os.kill(pid, signal.SIGKILL)
                        os.waitpid(pid, 0)
                    except (ProcessLookupError, ChildProcessError):
                        pass


def main():
    if os.environ.get("CTHULHU_NATIVE_REMOTE_KEYS") != "1":
        raise RuntimeError("Run only through native_remote_keys.py")
    # Adopt orphaned helpers/SDK peers after owner SIGKILL, without relying on
    # PID 1 or an init system to reap them. This process lives only in the fixture.
    libc = C.CDLL(None, use_errno=True)
    if libc.prctl(36, 1, 0, 0, 0) != 0:  # PR_SET_CHILD_SUBREAPER
        raise OSError(C.get_errno(), "Cannot adopt private fixture descendants")
    for mode in ("kill", "stop"):
        scenario(mode)


if __name__ == "__main__":
    main()