#!/bin/bash
# guest_mode.sh - read the NVIDIA X driver's nvClip screens[0].modeNum.
#
# The X driver (nvidia_drv.so) creates the nvClip shared area through its
# anonymous-file helper: memfd_create first (name derived from
# "/.nvidia_drv.XXXXXX"), then O_TMPFILE, then mkstemp+unlink.  570.144 size
# is 0x245000 bytes.  The fd is handed to GLX clients over X SEND_FDS.
# Layout (570.144, verified against nvidia_drv.so 570.144):
#   +0x00 u32 grabbed
#   +0x04 u32 token              (unique X-server token, sanity check)
#   +0x08 u32 screens[0].modeNum <-- this script
#   screens[] stride 0x84 (132) bytes
#
# modeNum parity is the GL client's __GL_DEV_MODESWITCH_TEST_COMPLETE bit:
#   even = device enabled / mode switch complete
#   odd  = mode switch in progress / device disabled
# GZDoom's GL init must observe an EVEN value; odd makes __glNVStateInit
# clear initBits, buffer objects get no GPU memory, and GZDoom crashes at
# 0x20 in FFlatVertexBuffer::Copy.
#
# STDOUT contract (parse this, nothing else):
#   MODENUM <n>
#   NO_XORG
#
# Exit codes:
#   0  MODENUM printed
#   3  NO_XORG (no Xorg process with an nvClip fd)
#   4  permission / unexpected error (NO_XORG still printed)
#
# Usage:
#   sudo ./guest_mode.sh                 # one shot, display from $DISPLAY (default :0)
#   sudo ./guest_mode.sh --wait 15       # poll up to 15 s for Xorg's nvClip fd
#   sudo ./guest_mode.sh --pid 1515      # force one Xorg pid (testing/forensics)
#   sudo ./guest_mode.sh --interval 50   # poll interval ms (with --wait)
#   sudo ./guest_mode.sh --debug         # candidate details on stderr
#
# Requires: bash, python3 (stdlib only), root (or a user that can read
# /proc/<Xorg>/fd; Xorg normally runs as root in the guest).

set -u

if ! command -v python3 >/dev/null 2>&1; then
    echo "guest_mode.sh: python3 not found" >&2
    echo "NO_XORG"
    exit 4
fi

exec python3 - "$@" <<'PY'
import argparse
import os
import re
import struct
import sys
import time

NVCLIP_TARGET_RE = re.compile(r'(?:^|/)\.nvidia_drv\.[^/]*$|memfd:.*nvidia_drv')
# 570.144 initial size is 0x245000.  ExtendSharedArea() ftruncates upward
# when a second channel block is needed, so accept a range, prefer exact.
NVCLIP_SIZE_EXACT = 0x245000
NVCLIP_SIZE_MIN = 0x80000
NVCLIP_SIZE_MAX = 0x800000
MODE_OFF = 0x8
TOKEN_OFF = 0x4


def eprint(msg):
    print(msg, file=sys.stderr)


def parse_args():
    p = argparse.ArgumentParser(add_help=True)
    p.add_argument('--pid', type=int, default=0,
                   help='only inspect this pid (bypasses the Xorg comm check)')
    p.add_argument('--display', default=os.environ.get('DISPLAY', ':0'),
                   help='X display to prefer (default $DISPLAY or :0)')
    p.add_argument('--wait', type=float, default=0.0,
                   help='poll up to SECONDS for the fd (default 0 = one shot)')
    p.add_argument('--interval', type=float, default=200.0,
                   help='poll interval in milliseconds (default 200)')
    p.add_argument('--debug', action='store_true',
                   help='print candidate details to stderr')
    return p.parse_args()


def read_text(path):
    with open(path, 'rb') as f:
        return f.read()


def proc_args(pid):
    try:
        raw = read_text('/proc/%d/cmdline' % pid).split(b'\0')
    except OSError:
        return []
    return [x.decode('utf-8', 'replace') for x in raw if x]


def proc_comm(pid):
    try:
        return read_text('/proc/%d/comm' % pid).decode('utf-8', 'replace').strip()
    except OSError:
        return ''


def proc_starttime(pid):
    # /proc/pid/stat field 22 (1-based).  Fields after the comm field:
    # rest[0] is field 3 (state), so starttime is rest[19].
    try:
        data = read_text('/proc/%d/stat' % pid).decode('utf-8', 'replace')
        rest = data.rsplit(')', 1)[1].split()
        return int(rest[19])
    except (OSError, IndexError, ValueError):
        return 0


def display_number(display):
    # ":0", ":0.0", "host:0.0" -> 0
    try:
        tail = display.rsplit(':', 1)[1]
        return int(tail.split('.', 1)[0])
    except (IndexError, ValueError):
        return None


def is_xorg(pid):
    comm = proc_comm(pid)
    if comm in ('Xorg', 'X'):
        return True
    args = proc_args(pid)
    if not args:
        return False
    base = os.path.basename(args[0])
    return base in ('Xorg', 'X')


def display_matches(pid, disp):
    if disp is None:
        return False
    needle = ':%d' % disp
    for arg in proc_args(pid):
        if arg == needle or arg.startswith(needle + '.'):
            return True
    return False


def candidate_fds(pid):
    """Yield (fd_name, target, size) for nvClip-like fds of pid."""
    fd_dir = '/proc/%d/fd' % pid
    try:
        names = os.listdir(fd_dir)
    except PermissionError:
        raise
    except OSError:
        return
    for name in names:
        if not name.isdigit():
            continue
        path = os.path.join(fd_dir, name)
        try:
            target = os.readlink(path)
            st = os.stat(path)
        except PermissionError:
            raise
        except OSError:
            continue
        if not NVCLIP_TARGET_RE.search(target):
            continue
        if not (NVCLIP_SIZE_MIN <= st.st_size <= NVCLIP_SIZE_MAX):
            continue
        yield name, target, st.st_size


def read_candidate(pid, fd_name, target, size):
    path = '/proc/%d/fd/%s' % (pid, fd_name)
    try:
        fd = os.open(path, os.O_RDONLY)
    except OSError as exc:
        return None, str(exc)
    try:
        blob = os.pread(fd, 12, 0)
    except OSError as exc:
        return None, str(exc)
    finally:
        os.close(fd)
    if len(blob) < 12:
        return None, 'short read (%d bytes)' % len(blob)
    grabbed, token, mode = struct.unpack_from('<III', blob, 0)
    if mode > 0xFFFF:
        return None, 'implausible modeNum 0x%x' % mode
    return {
        'pid': pid,
        'fd': int(fd_name),
        'target': target,
        'size': size,
        'grabbed': grabbed,
        'token': token,
        'mode': mode,
        'start': proc_starttime(pid),
        'disp': display_matches(pid, display_number(args.display)),
    }, None


def scan(debug=False):
    """Return (best_candidate_or_None, saw_permission_error)."""
    best = None
    perm_error = False

    try:
        pids = [int(x) for x in os.listdir('/proc') if x.isdigit()]
    except OSError as exc:
        eprint('guest_mode.sh: cannot list /proc: %s' % exc)
        return None, False

    for pid in pids:
        if args.pid and pid != args.pid:
            continue
        if not args.pid and not is_xorg(pid):
            continue
        try:
            fds = list(candidate_fds(pid))
        except PermissionError:
            perm_error = True
            if debug:
                eprint('pid %d: permission denied listing fd/' % pid)
            continue
        for fd_name, target, size in fds:
            cand, err = read_candidate(pid, fd_name, target, size)
            if cand is None:
                if debug:
                    eprint('pid %d fd %s: %s' % (pid, fd_name, err))
                continue
            if debug:
                eprint('candidate pid=%d fd=%s size=0x%x token=0x%x mode=%d '
                       'start=%d display_match=%s target=%s' %
                       (cand['pid'], cand['fd'], cand['size'], cand['token'],
                        cand['mode'], cand['start'], cand['disp'], cand['target']))
            # Score: display match, exact size, nonzero token, newest Xorg.
            key = (cand['disp'],
                   1 if cand['size'] == NVCLIP_SIZE_EXACT else 0,
                   1 if cand['token'] != 0 else 0,
                   cand['start'],
                   -cand['pid'])
            if best is None or key > best[0]:
                best = (key, cand)

    return (best[1] if best else None), perm_error


def main():
    global args
    args = parse_args()
    deadline = time.monotonic() + max(0.0, args.wait)
    interval = max(0.01, args.interval / 1000.0)
    perm_error = False

    while True:
        cand, pe = scan(args.debug)
        perm_error = perm_error or pe
        if cand is not None:
            print('MODENUM %d' % cand['mode'])
            return 0
        if time.monotonic() >= deadline:
            break
        time.sleep(interval)

    if perm_error:
        eprint('guest_mode.sh: Xorg exists but /proc/<pid>/fd is not readable; '
               'run as root (sudo).')
        print('NO_XORG')
        return 4
    print('NO_XORG')
    return 3


if __name__ == '__main__':
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print('NO_XORG')
        sys.exit(4)
    except Exception as exc:  # keep the stdout contract even on bugs
        eprint('guest_mode.sh: unexpected error: %s' % exc)
        print('NO_XORG')
        sys.exit(4)
PY