#!/usr/bin/env python3
"""Box supervisor. Exec preserves the PID expected by kiosk window activation.

Metadata travels over pinned-host-key SSH with a per-host forced-command key.
The private state directory is not a public socket/port. Native cursor resources
are created and freed by Moonlight's SDL main thread, never by this process.
"""
import base64
import hashlib
import json
import os
from pathlib import Path
import select
import shutil
import signal
import struct
import subprocess
import sys
import tempfile
import time
import uuid

HEADER = struct.Struct('<8sQQQIIIIiiII')
MAX_LINE = 100000
STATE_TIMEOUT = 2.0


def verify_client(path, expected):
    if (not isinstance(expected, str) or len(expected) != 64 or
            hashlib.sha256(Path(path).read_bytes()).hexdigest() != expected):
        raise ValueError('native cursor client build changed; verify build before enabling profile')


def validate(state, lease, size, previous):
    if type(state.get('version')) is not int or state.get('version') != 1 or state.get('lease') != lease or state.get('type') != 'cursor':
        raise ValueError('cursor protocol/session mismatch')
    fields = ('seq', 'ack', 'w', 'h', 'hx', 'hy', 'x', 'y', 'sw', 'sh')
    if any(type(state.get(k)) is not int for k in fields):
        raise ValueError('noninteger field')
    if not previous < state['seq'] < 2**63:
        raise ValueError('stale sequence')
    if not 0 <= state['ack'] < 2**63:
        raise ValueError('invalid input acknowledgement')
    w, h = state['w'], state['h']
    if not (0 < w <= 128 and 0 < h <= 128 and 0 <= state['hx'] < w and 0 <= state['hy'] < h):
        raise ValueError('cursor bounds')
    if [state['sw'], state['sh']] != size:
        raise ValueError('display changed')
    if not (0 <= state['x'] < size[0] and 0 <= state['y'] < size[1]):
        raise ValueError('cursor position outside single display')
    pixels = base64.b64decode(state['rgba'], validate=True)
    if len(pixels) != w * h * 4:
        raise ValueError('pixel length')
    return HEADER.pack(b'VDICUR1\0', int(lease[:16], 16), state['seq'],
                       time.monotonic_ns(), w, h, state['hx'], state['hy'],
                       state['x'], state['y'], *size) + pixels + struct.pack('<Q', state['ack'])


def write_state(directory, payload):
    tmp = directory / 'state.tmp'
    with tmp.open('wb') as f:
        f.write(payload)
    tmp.replace(directory / 'state')


def supervise(profile, size, lease, directory, ready_fd, parent):
    ssh = None
    ready = False
    try:
        cmd = ['ssh', '-T', '-i', profile['key'], '-o', 'IdentitiesOnly=yes',
               '-o', 'BatchMode=yes', '-o', 'StrictHostKeyChecking=yes',
               '-o', 'UserKnownHostsFile=' + profile['known_hosts'],
               '-o', 'ConnectTimeout=5', '-o', 'ServerAliveInterval=1',
               '-o', 'ServerAliveCountMax=2', profile['user'] + '@' + profile['host']]
        ssh = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        ssh.stdin.write((json.dumps(dict(version=1, lease=lease, size=size)) + '\n').encode())
        ssh.stdin.flush()
        started = last = last_ping = time.monotonic()
        seq, buffer, acknowledged = 0, b'', False
        last_input = 0
        while True:
            # A kiosk SIGKILL of the exec'd client must still tear down the lease.
            if os.getppid() != parent:
                break
            now = time.monotonic()
            if now - last_ping >= .25:
                ssh.stdin.write(b'ping\n')
                ssh.stdin.flush()
                last_ping = now
            input_path = directory / 'state.input'
            if ready and input_path.exists():
                data = input_path.read_bytes()
                if len(data) != 16:
                    raise ValueError('invalid local input record')
                input_seq, px, py = struct.unpack('<Qii', data)
                if input_seq > last_input:
                    if not (input_seq < 2**63 and 0 <= px < size[0] and 0 <= py < size[1]):
                        raise ValueError('local input bounds')
                    ssh.stdin.write((json.dumps(dict(input=input_seq, x=px, y=py))+'\n').encode())
                    ssh.stdin.flush()
                    last_input = input_seq
            if (ready and now - last > STATE_TIMEOUT) or (not ready and now - started > 50):
                raise RuntimeError('native cursor heartbeat timed out (%.3fs)' % (now - last))
            if select.select([ssh.stdout], [], [], .025)[0]:
                chunk = os.read(ssh.stdout.fileno(), 65536)
                if not chunk:
                    raise RuntimeError('native cursor SSH ended')
                buffer += chunk
                if len(buffer) > MAX_LINE * 2:
                    raise ValueError('cursor receive buffer limit')
                while b'\n' in buffer:
                    line, buffer = buffer.split(b'\n', 1)
                    if len(line) > MAX_LINE:
                        raise ValueError('cursor packet limit')
                    state = json.loads(line)
                    if not acknowledged:
                        if state != dict(type='ready', version=1, lease=lease):
                            raise ValueError('cursor handshake rejected')
                        acknowledged = True
                        continue
                    payload = validate(state, lease, size, seq)
                    seq = state['seq']
                    write_state(directory, payload)
                    last = time.monotonic()
                    if not ready:
                        os.write(ready_fd, b'ready\n')
                        os.close(ready_fd)
                        ready = True
                        print('VDI native cursor: authenticated cursor-free session ready',
                              file=sys.stderr, flush=True)
        return 0
    except Exception as error:
        print('VDI native cursor: ' + str(error), file=sys.stderr, flush=True)
        if ready and os.getppid() == parent:
            os.kill(parent, signal.SIGTERM)
        return 78
    finally:
        if ssh:
            try:
                ssh.stdin.write(b'stop\n')
                ssh.stdin.flush()
                ssh.stdin.close()
                ssh.wait(timeout=12)
            except (BrokenPipeError, subprocess.TimeoutExpired, OSError):
                ssh.terminate()
                try:
                    ssh.wait(timeout=3)
                except subprocess.TimeoutExpired:
                    ssh.kill()
        shutil.rmtree(directory)


def launch(argv):
    client, args = argv[0], argv[1:]
    if not (len(args) >= 2 and args[0] == 'stream'):
        os.execv(client, argv)
    host = args[1]
    if any(c not in '0123456789.' for c in host):
        os.execv(client, argv)
    config = Path.home() / '.config/vdi-terminal/native-cursor' / (host + '.json')
    if os.environ.get('VDI_NATIVE_CURSOR') == '0' or not config.exists():
        os.execv(client, argv)
    profile = json.loads(config.read_text())
    # An old or separately rebuilt client must never enter cursor-free video
    # without the SDL receiver. Verify BEFORE acquiring any remote lease.
    verify_client(client, profile.get('client_sha256'))
    if profile['host'] != host or '--absolute-mouse' not in args:
        raise ValueError('native cursor needs matching absolute desktop mode')
    resolution = args[args.index('--resolution') + 1]
    size = list(map(int, resolution.split('x')))
    if size != profile['size']:
        raise ValueError('native cursor needs configured single-screen size')
    lease = uuid.uuid4().hex
    directory = Path(tempfile.mkdtemp(prefix='vdi-native-cursor-',
                                    dir=os.environ.get('XDG_RUNTIME_DIR', '/tmp')))
    rfd, wfd = os.pipe()
    parent = os.getpid()
    worker = os.fork()
    if worker == 0:
        os.close(rfd)
        os._exit(supervise(profile, size, lease, directory, wfd, parent))
    os.close(wfd)
    if not select.select([rfd], [], [], 55)[0] or os.read(rfd, 32) != b'ready\n':
        # No normal-mode launch while a cursor-free producer might still live.
        raise RuntimeError('native cursor preflight failed; reconnect after recovery')
    os.close(rfd)
    os.environ['VDI_NATIVE_CURSOR_STATE'] = str(directory / 'state')
    os.environ['VDI_NATIVE_CURSOR_EPOCH'] = str(int(lease[:16], 16))
    os.execv(client, argv)


if __name__ == '__main__':
    launch(sys.argv[1:])