#!/usr/bin/env python3
"""Authenticated single-seat Linux cursor lease; run by a restricted SSH key.

The guard owns the Sunshine lifecycle. Cursor-free capture starts BEFORE the
new video connection; ending the lease stops that capture process before
restoring normal capture. There is deliberately no mid-stream mode switch.
"""
import base64
import ctypes as C
import fcntl
import hashlib
import json
import os
from pathlib import Path
import select
import signal
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent
RUN = Path('/run/user/%d/vdi-native-cursor' % os.getuid())
DROP = RUN.parent / 'systemd/user/vdi-sunshine.service.d/90-native-cursor.conf'
SERVICE = 'vdi-sunshine.service'
GUARD = 'vdi-native-cursor-guard.service'
TTL = 4.0


def systemctl(*args):
    return subprocess.run(['systemctl', '--user'] + list(args), check=True,
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                          timeout=20).stdout.decode().strip()


def recover():
    # Stop the cursor-free producer FIRST. A new normal producer cannot resume
    # an old Moonlight connection, so lost SSH never creates a double cursor.
    if DROP.exists():
        systemctl('stop', SERVICE)
        DROP.unlink()
        systemctl('daemon-reload')
        systemctl('start', SERVICE)
    for name in ('lease', 'ready', 'receipt'):
        (RUN / name).unlink(missing_ok=True)


def guard():
    lease = (RUN / 'lease').read_text().strip()
    if len(lease) != 32 or any(c not in '0123456789abcdef' for c in lease):
        raise ValueError('invalid lease')
    manifest = json.loads((ROOT / 'manifest.json').read_text())
    for name, expected in manifest.items():
        if name not in ('sunshine', 'capture-no-cursor.so') or hashlib.sha256(
                (ROOT / name).read_bytes()).hexdigest() != expected:
            raise RuntimeError('pinned capture build changed')
    kms_config = ROOT / 'kms.json'
    scanout = None
    def checked_scanout(timeout=3):
        deadline = time.monotonic() + timeout
        last_error = None
        while True:
            try:
                return scanout.check()
            except RuntimeError as exc:
                last_error = exc
                if time.monotonic() >= deadline:
                    raise last_error
                time.sleep(.1)
    if kms_config.exists():
        from kms import VirtioScanout
        scanout = VirtioScanout(json.loads(kms_config.read_text()))
        # Wake an idle output before inspecting its active scanout. This does
        # not unlock the session or disable the user's screen lock.
        subprocess.run(['/usr/bin/xset', 'dpms', 'force', 'on'], check=True,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=5)
        if hashlib.sha256(Path('/usr/bin/sunshine').read_bytes()).hexdigest() != manifest['sunshine']:
            raise RuntimeError('pinned system capture build changed')
    systemctl('stop', SERVICE)
    DROP.parent.mkdir(parents=True, exist_ok=True)
    if scanout:
        # File capabilities remain on the original, verified system binary.
        # Command-line override is lease-local; recover restores X11 capture.
        DROP.write_text('[Service]\nExecStart=\nExecStart=/usr/bin/sunshine %s capture=kms\n' %
                        (Path.home() / '.config/sunshine/sunshine.conf'))
    else:
        DROP.write_text('[Service]\nEnvironment="LD_PRELOAD=%s" "VDI_CURSOR_LEASE=%s" '
                    '"VDI_CURSOR_RECEIPT=%s"\n' %
                    (ROOT / 'capture-no-cursor.so', lease, RUN / 'receipt') +
                    'ExecStart=\nExecStart=%s %s\n' %
                    (ROOT / 'sunshine', Path.home() / '.config/sunshine/sunshine.conf'))
    systemctl('daemon-reload')
    # The normal X11 producer can leave the DRM plane idle while it is being
    # stopped. Install the recoverable drop-in first, then validate the
    # stopped display before launching KMS. ExecStopPost restores X11 if this
    # fail-closed check rejects the current output.
    if scanout:
        subprocess.run(['/usr/bin/xset', 'dpms', 'force', 'on'], check=True,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=5)
        checked_scanout()
    systemctl('start', SERVICE)
    # dlsym receipt proves this exact process resolved the adapter entrypoint.
    deadline = time.monotonic() + 12
    while not scanout and not (RUN / 'receipt').exists():
        if time.monotonic() > deadline:
            raise RuntimeError('capture adapter not resolved')
        time.sleep(.05)
    pid = int(systemctl('show', SERVICE, '-p', 'MainPID', '--value'))
    if scanout:
        checked_scanout()
        if b'capture=kms\0' not in Path('/proc/%d/cmdline' % pid).read_bytes():
            raise RuntimeError('KMS capture override not active')
    elif str(ROOT / 'capture-no-cursor.so') not in Path('/proc/%d/maps' % pid).read_text():
        raise RuntimeError('capture adapter not mapped')
    (RUN / 'ready').write_text(lease)
    while (RUN / 'lease').exists():
        if time.time() - (RUN / 'lease').stat().st_mtime > TTL:
            break
        if int(systemctl('show', SERVICE, '-p', 'MainPID', '--value')) != pid:
            break  # Never accept an encoder restart inside the same lease.
        if scanout:
            checked_scanout(1)
        time.sleep(.15)


class Cursor(C.Structure):
    _fields_ = [('x', C.c_short), ('y', C.c_short), ('w', C.c_ushort),
                ('h', C.c_ushort), ('hx', C.c_ushort), ('hy', C.c_ushort),
                ('serial', C.c_ulong), ('pixels', C.POINTER(C.c_ulong)),
                ('atom', C.c_ulong), ('name', C.c_char_p)]


class XCursor:
    def __init__(self):
        # Absolute original library path; never inherit the capture adapter.
        self.x = C.CDLL('libX11.so.6')
        self.f = C.CDLL('/lib/x86_64-linux-gnu/libXfixes.so.3')
        self.x.XOpenDisplay.argtypes = [C.c_char_p]
        self.x.XOpenDisplay.restype = C.c_void_p
        self.x.XDefaultScreen.argtypes = [C.c_void_p]
        self.x.XDefaultRootWindow.argtypes = [C.c_void_p]
        self.x.XDefaultRootWindow.restype = C.c_ulong
        self.x.XGetGeometry.argtypes = [C.c_void_p, C.c_ulong, C.POINTER(C.c_ulong),
            C.POINTER(C.c_int), C.POINTER(C.c_int), C.POINTER(C.c_uint),
            C.POINTER(C.c_uint), C.POINTER(C.c_uint), C.POINTER(C.c_uint)]
        for name in ('XDisplayWidth', 'XDisplayHeight'):
            getattr(self.x, name).argtypes = [C.c_void_p, C.c_int]
        self.x.XFree.argtypes = [C.c_void_p]
        self.f.XFixesGetCursorImage.argtypes = [C.c_void_p]
        self.f.XFixesGetCursorImage.restype = C.POINTER(Cursor)
        self.display = self.x.XOpenDisplay(None)
        if not self.display:
            raise RuntimeError('no authenticated X11 display')
        self.screen = self.x.XDefaultScreen(self.display)
        self.root = self.x.XDefaultRootWindow(self.display)
        self.size = self.dimensions()

    def dimensions(self):
        root = C.c_ulong()
        x, y = C.c_int(), C.c_int()
        w, h, border, depth = [C.c_uint() for _ in range(4)]
        if not self.x.XGetGeometry(self.display, self.root, C.byref(root),
                C.byref(x), C.byref(y), C.byref(w), C.byref(h),
                C.byref(border), C.byref(depth)):
            raise RuntimeError('root geometry unavailable')
        return w.value, h.value

    def read(self):
        if self.dimensions() != self.size:
            raise RuntimeError('display changed; reconnect required')
        ptr = self.f.XFixesGetCursorImage(self.display)
        if not ptr:
            raise RuntimeError('cursor unavailable')
        try:
            c = ptr.contents
            if not (0 <= c.w <= 128 and 0 <= c.h <= 128):
                raise ValueError('unsupported cursor bounds')
            rgba = bytearray()
            for i in range(c.w * c.h):
                p = c.pixels[i]
                a = (p >> 24) & 255
                # XRender pixels are premultiplied; SDL expects straight RGBA.
                rgb = [min(255, (((p >> s) & 255) * 255 + a // 2) // a)
                       if a else 0 for s in (16, 8, 0)]
                rgba.extend(rgb + [a])
            w, h, hx, hy = c.w, c.h, c.hx, c.hy
            if not any(rgba[3::4]):
                # GTK's blank cursor can be 1x1 with a hotspot of (1,1).
                # Canonicalize invisible pixels; never clip a visible hotspot.
                w, h, hx, hy, rgba = 1, 1, 0, 0, bytearray(4)
            elif not (hx < w and hy < h):
                raise ValueError('visible cursor hotspot outside image')
            return dict(w=w, h=h, hx=hx, hy=hy, x=c.x, y=c.y,
                        sw=self.size[0], sh=self.size[1], shape=int(c.serial),
                        rgba=base64.b64encode(rgba).decode())
        finally:
            self.x.XFree(ptr)


def stream():
    RUN.mkdir(mode=0o700, parents=True, exist_ok=True)
    with (RUN / 'lock').open('w') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        # Version/capabilities come through the authenticated SSH stdin.
        if not select.select([sys.stdin], [], [], 5)[0]:
            raise RuntimeError('missing hello')
        hello = json.loads(sys.stdin.buffer.readline(2048))
        lease = hello.get('lease', '')
        if hello.get('version') != 1 or len(lease) != 32 or any(
                c not in '0123456789abcdef' for c in lease):
            raise ValueError('bad hello')
        x = XCursor()
        if list(x.size) != hello.get('size'):
            raise ValueError('source size differs')
        # Fail before touching the service if the native shape is unsupported.
        x.read()
        systemctl('stop', GUARD)
        (RUN / 'lease').write_text(lease)
        systemctl('start', GUARD)
        deadline = time.monotonic() + 40
        try:
            while not (RUN / 'ready').exists():
                if time.monotonic() > deadline:
                    raise RuntimeError('guard not ready')
                os.utime(RUN / 'lease', None)
                time.sleep(.05)
            if (RUN / 'ready').read_text() != lease:
                raise RuntimeError('stale guard')
            print(json.dumps(dict(type='ready', version=1, lease=lease)), flush=True)
            seq = 0
            last_shape = None
            last_sent = 0
            heartbeat = time.monotonic()
            acknowledged_input = 0
            pending_input = None
            while True:
                now = time.monotonic()
                if select.select([sys.stdin], [], [], 0)[0]:
                    line = sys.stdin.buffer.readline(2048)
                    if not line or line.strip() == b'stop':
                        break
                    if line.strip() != b'ping':
                        request = json.loads(line)
                        if (set(request) != {'input', 'x', 'y'} or
                            any(type(request[k]) is not int for k in request) or
                            not acknowledged_input < request['input'] < 2**63 or
                            not 0 <= request['x'] < x.size[0] or
                            not 0 <= request['y'] < x.size[1]):
                            raise ValueError('bad input acknowledgement request')
                        pending_input = request
                    heartbeat = now
                if now - heartbeat > TTL:
                    raise RuntimeError('client lease expired')
                os.utime(RUN / 'lease', None)
                if not (RUN / 'ready').exists():
                    raise RuntimeError('capture guard ended')
                read_started = time.monotonic()
                state = x.read()
                read_seconds = time.monotonic() - read_started
                if read_seconds > .1:
                    print('VDI native cursor: X11 read blocked %.3fs' % read_seconds,
                          file=sys.stderr, flush=True)
                if pending_input and (abs(state['x'] - pending_input['x']) <= 2 and
                                      abs(state['y'] - pending_input['y']) <= 2):
                    acknowledged_input = pending_input['input']
                    pending_input = None
                state['ack'] = acknowledged_input
                identity = (state['shape'], state['rgba'])
                if identity != last_shape or now - last_sent >= .05:
                    seq += 1
                    state.update(type='cursor', version=1, lease=lease, seq=seq)
                    print(json.dumps(state, separators=(',', ':')), flush=True)
                    last_shape, last_sent = identity, now
                time.sleep(.008)
        finally:
            (RUN / 'lease').unlink(missing_ok=True)
            systemctl('stop', GUARD)


def main():
    os.environ['DISPLAY'] = ':0'
    os.environ['XDG_RUNTIME_DIR'] = '/run/user/%d' % os.getuid()
    os.environ['DBUS_SESSION_BUS_ADDRESS'] = 'unix:path=%s/bus' % os.environ['XDG_RUNTIME_DIR']
    gdm = Path(os.environ['XDG_RUNTIME_DIR']) / 'gdm/Xauthority'
    os.environ['XAUTHORITY'] = str(gdm if gdm.exists() else Path.home() / '.Xauthority')
    for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal.signal(sig, lambda *_: sys.exit(0))
    mode = sys.argv[1] if len(sys.argv) > 1 else 'stream'
    if mode == 'recover':
        recover()
    elif mode == 'guard':
        guard()
    elif mode == 'stream':
        stream()
    else:
        raise ValueError('unknown mode')


if __name__ == '__main__':
    main()