#!/usr/bin/env python3
"""Opt-in root acceptance test: private Xvfb and disposable PAM user, never live X.

Usage: sudo python3 scripts/verify-lockscreen --backend /path/to/monocoat-locker --out /tmp/lock-proof
Optional --image /path/to/prepared.png --avatar /path/to/avatar.png --fonts /path/to/fonts.
"""
import argparse
import ctypes
import hashlib
import os
import pwd
import secrets
import select
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
from monocoat.lockscreen import arguments  # noqa: E402


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--backend", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--image", type=Path)
    parser.add_argument("--avatar", type=Path)
    parser.add_argument("--glass", type=Path)
    parser.add_argument("--fonts", type=Path)
    parser.add_argument("--submit", choices=("click", "space"), default="click")
    args = parser.parse_args()
    if os.geteuid() != 0:
        parser.error("root is required to create and remove the disposable PAM account")
    args.out.mkdir(parents=True, exist_ok=True)
    username = f"mc-lock-test-{os.getpid()}"
    password = "Mc-" + secrets.token_hex(16)
    with tempfile.TemporaryDirectory(prefix="monocoat-lock-test-") as temporary:
        home = Path(temporary)
        home.chmod(0o755)
        subprocess.run(["useradd", "--system", "--no-create-home", "--home-dir", str(home),
                        "--shell", "/bin/sh", username], check=True)
        xserver = locker = None
        try:
            subprocess.run(["chpasswd"], input=f"{username}:{password}\n", text=True, check=True)
            user = pwd.getpwnam(username)
            os.chown(home, user.pw_uid, user.pw_gid)
            if args.fonts:
                shutil.copytree(args.fonts, home / "fonts")
            fontconfig = home / "fonts.conf"
            fontconfig.write_text(f"<fontconfig><dir>{home}/fonts</dir><dir>/usr/share/fonts</dir><cachedir>{home}/fontcache</cachedir></fontconfig>")
            xserver = subprocess.Popen(["Xvfb", "-displayfd", "1", "-screen", "0", "1920x1080x24",
                                        "-dpi", "96", "-ac", "-nolisten", "tcp"],
                                       stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
            display = ":" + xserver.stdout.readline().strip()
            env = dict(os.environ, HOME=str(home), DISPLAY=display, XAUTHORITY="/dev/null",
                       XDG_CONFIG_HOME=str(home / "config"), XDG_DATA_HOME=str(home / "data"),
                       XDG_CACHE_HOME=str(home / "cache"), FONTCONFIG_FILE=str(fontconfig),
                       DBUS_SESSION_BUS_ADDRESS=f"unix:path={home}/no-bus", NO_AT_BRIDGE="1",
                       MONOCOAT_LOCK_STYLE="macos")
            # Art may live in the real user's private cache. Copy it into the
            # disposable account's sandbox instead of silently testing fallback.
            image = None
            for source, name, variable in ((args.image, "wallpaper.png", None),
                                          (args.avatar, "avatar.png", "MONOCOAT_LOCK_AVATAR"),
                                          (args.glass, "glass.png", "MONOCOAT_LOCK_GLASS")):
                if source:
                    target = home / name
                    shutil.copyfile(source, target)
                    target.chmod(0o644)
                    if variable:
                        env[variable] = str(target)
                    else:
                        image = target
            read_fd, write_fd = os.pipe()
            env["XSS_SLEEP_LOCK_FD"] = str(write_fd)
            argv = arguments(image, "Blake Bratcher", nofork=True)
            argv[0] = str(args.backend.resolve())
            locker = subprocess.Popen(argv, env=env, user=user.pw_uid, group=user.pw_gid,
                                      extra_groups=[], start_new_session=True, pass_fds=(write_fd,),
                                      stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            os.close(write_fd)
            try:
                assert select.select([read_fd], [], [], 25)[0], "sleep inhibitor never released"
                assert os.read(read_fd, 1) == b"", "invalid inhibitor state"
            finally:
                os.close(read_fd)
            assert locker.poll() is None, "locker failed before input grab"
            xlib = ctypes.CDLL("libX11.so.6")
            xlib.XOpenDisplay.argtypes = [ctypes.c_char_p]
            xlib.XOpenDisplay.restype = ctypes.c_void_p
            connection = xlib.XOpenDisplay(display.encode())
            xlib.XDefaultRootWindow.argtypes = [ctypes.c_void_p]
            xlib.XDefaultRootWindow.restype = ctypes.c_ulong
            root = xlib.XDefaultRootWindow(connection)
            xlib.XGrabKeyboard.argtypes = [ctypes.c_void_p, ctypes.c_ulong, ctypes.c_int,
                                          ctypes.c_int, ctypes.c_int, ctypes.c_ulong]
            assert xlib.XGrabKeyboard(connection, root, 0, 1, 1, 0) == 1, "keyboard is not grabbed"
            xlib.XCloseDisplay.argtypes = [ctypes.c_void_p]
            xlib.XCloseDisplay(connection)
            def key(*keys, clear=True):
                subprocess.run(["xdotool", "key", *(["--clearmodifiers"] if clear else []), *keys], env=env, check=True)
            def capture(name):
                time.sleep(0.25)  # allow the locker to process another X client's events
                subprocess.run(["import", "-limit", "thread", "2", "-define", "png:compression-level=1", "-display", display, "-window", "root",
                                str(args.out / name)], check=True)
            def pixels(crop):
                # Pixel bytes exclude PNG timestamps and the ticking clock.
                raw = subprocess.check_output(["import", "-limit", "thread", "2", "-display", display, "-window", "root",
                                               "-crop", crop, "-depth", "8", "rgb:-"])
                return hashlib.sha256(raw).hexdigest()
            def mouse(*commands):
                subprocess.run(["xdotool", *commands], env=env, check=True)
            time.sleep(0.5)
            def cpu_seconds():
                fields = Path(f"/proc/{locker.pid}/stat").read_text().rsplit(")", 1)[1].split()
                return (int(fields[11]) + int(fields[12])) / os.sysconf("SC_CLK_TCK")
            cpu_before = cpu_seconds()
            time.sleep(1.2)
            print(f"Idle redraw: {(cpu_seconds() - cpu_before)*1000:.0f} ms CPU / 1.2 s", flush=True)
            capture("lockscreen.png")
            # Empty arrow is inert. Help opens on release, and typing/ESC
            # dismisses it without weakening the keyboard/pointer grab.
            mouse("mousemove", "1051", "960", "click", "1")
            assert locker.poll() is None, "empty arrow unlocked"
            plain_help = pixels("300x46+810+992")
            mouse("mousemove", "1096", "960")
            time.sleep(0.2)
            mouse("click", "1")
            capture("password-help.png")
            mouse_help = pixels("300x46+810+992")
            assert mouse_help != plain_help, "clicking help did not open its popover"
            key("Escape")
            time.sleep(0.2)
            normal_hint = pixels("300x20+810+1003")
            key("Caps_Lock", clear=False)
            time.sleep(0.1)
            capture("caps-lock.png")
            assert pixels("300x20+810+1003") != normal_hint, "Caps Lock did not update the hint"
            key("Caps_Lock", clear=False)
            # With no password, Tab skips the disabled submit and reaches Help.
            key("Tab", "Return")
            capture("keyboard-help.png")
            assert pixels("300x46+810+992") == mouse_help, "keyboard help did not match mouse help"
            key("Escape")
            key("Escape", "ctrl+c")
            assert locker.poll() is None, "Escape or Ctrl+C bypassed the lock"
            key("Escape")
            subprocess.run(["xdotool", "type", "--clearmodifiers", "incorrect-password"], env=env, check=True)
            mouse("mousemove", "1051", "960")
            time.sleep(0.2)
            capture("submit-hover.png")
            # A drag away cancels the click; auth must not run on mouse-down.
            mouse("mousedown", "1")
            capture("submit-pressed.png")
            assert locker.poll() is None
            mouse("mousemove", "1200", "900", "mouseup", "1")
            capture("password-masked.png")
            key("Return")
            time.sleep(3)
            assert locker.poll() is None, "incorrect password unlocked the session"
            capture("incorrect-password.png")
            time.sleep(3)
            key("Escape")
            subprocess.run(["xdotool", "type", "--clearmodifiers", password], env=env, check=True)
            # Shift+Tab must return from Submit to the field without adding
            # a tab or a space to the password. Reaching the real unlock later
            # proves navigation preserved the password buffer byte-for-byte.
            key("Tab", "shift+Tab")
            mouse("mousemove", "1051", "960", "mousedown", "1")
            time.sleep(0.1)
            assert locker.poll() is None, "mouse-down submitted a valid password"
            mouse("mousemove", "1200", "900", "mouseup", "1")
            time.sleep(0.1)
            assert locker.poll() is None, "dragging away submitted a valid password"
            mouse("mousemove", "950", "960", "click", "1")
            time.sleep(0.2)  # finish hover feedback before measuring caret edges
            samples = []
            for _ in range(7):
                samples.append(pixels("330x120+805+925"))
                time.sleep(0.12)
            assert len(set(samples)) >= 2, "focused caret never blinks"
            # The visible arrow must submit through the same upstream PAM path.
            if args.submit == "space":
                key("Tab", "space")
            else:
                subprocess.run(["xdotool", "mousemove", "1051", "960", "click", "1"], env=env, check=True)
            assert locker.wait(timeout=15) == 0, "valid password did not unlock"
            print(f"PASS: input grabbed before inhibitor release; empty/dragged arrows and Escape/Ctrl+C cannot unlock; help, Caps Lock and caret verified; Tab/Shift+Tab preserve password; wrong password rejected; valid password + {args.submit} unlocks.")
        finally:
            if locker:
                try:
                    os.killpg(locker.pid, signal.SIGTERM)
                except ProcessLookupError:
                    pass
                locker.wait()
            if xserver:
                xserver.terminate()
                xserver.wait()
            subprocess.run(["userdel", username], check=True)


if __name__ == "__main__":
    main()
