XFESystem

"""Real private-QEMU USB HID shortcuts; reuse one verified image without copying it.

This verifies the guest xHCI/USB boot keyboard path, not a physical Corsair
keyboard. No host keyboard, physical disk, or existing QEMU process is touched.
"""
import argparse
import hashlib
import json
import os
import pathlib
import re
import shutil
import socket
import subprocess
import time
import uuid

from ReadImage import inspect_image


ROOT = pathlib.Path(__file__).resolve().parents[2]
FAILURES = ('CPU EXCEPTION:', '[FAILFAST]', '*** UNHANDLED EXCEPTION',
            'XFE_RUNTIME_OOM|', 'XFE_RUNTIME_STACK_OOM|')
HEADER = re.compile(r'^XFE_DESKTOP_FLAGS\|([01])\|([01])\|Active\|(\d+)\r?\n', re.M)
ROW = re.compile(r'XFE_WINDOW_STATE\|(\d+)\|([^|]+)\|(-?\d+)\|(-?\d+)\|(\d+)\|(\d+)\|([01])\|([01])\|([01])')


def sha256(path):
    digest = hashlib.sha256()
    with pathlib.Path(path).open('rb') as source:
        for block in iter(lambda: source.read(1024 * 1024), b''):
            digest.update(block)
    return digest.hexdigest()


def snapshots(log):
    """Require a complete FLAGS/appearance/rows report and its serial seal."""
    headers = list(HEADER.finditer(log))
    result = []
    for index, header in enumerate(headers):
        end = headers[index + 1].start() if index + 1 < len(headers) else len(log)
        section = log[header.end():end]
        windows, appearance, sealed, malformed = {}, False, index + 1 < len(headers), False
        for line in section.splitlines(keepends=True):
            if not line.endswith('\n'):
                break
            line = line.rstrip('\r\n')
            if not appearance:
                appearance = bool(re.fullmatch(r'XFE_DESKTOP_APPEARANCE\|(Light|Dark)\|[0-9A-F]{6}\|Persistent\|[01]', line))
                if not appearance:
                    malformed = True
                    break
                continue
            row = ROW.fullmatch(line)
            if row is None:
                sealed = True
                break
            value = row.groups()
            identifier = int(value[0])
            if identifier in windows:
                malformed = True
                break
            windows[identifier] = {'Application': value[1], 'Bounds': [int(part) for part in value[2:6]],
                                   'Minimized': int(value[6]), 'Maximized': int(value[7]), 'Snapped': int(value[8])}
        active = int(header[3])
        if appearance and sealed and not malformed and (active == 0 or active in windows):
            result.append({'StartMenuOpen': int(header[1]), 'WindowSwitching': int(header[2]),
                           'ActiveId': active, 'Windows': windows, 'SerialOffset': header.start()})
    return result


def require_usb_only_topology(tree):
    assert 'dev: qemu-xhci, id "XfeUsbXhci"' in tree, 'The owned xHCI controller is absent'
    assert 'bus: XfeUsbXhci.0' in tree, 'The owned xHCI USB bus is absent'
    assert 'dev: usb-kbd, id "XfeUsbKeyboard"' in tree, 'The owned USB HID keyboard is absent'
    assert 'i8042' not in tree and 'ps2-kbd' not in tree, 'PS/2 input could mask a broken USB path'
    assert tree.count('dev: usb-kbd,') == 1, 'More than one USB keyboard makes input routing ambiguous'


def require_image_backing(blocks, image):
    expected = pathlib.Path(image).resolve()
    matches = []
    for block in blocks:
        node = block.get('inserted', {}).get('image', {})
        while node:
            if node.get('filename') and pathlib.Path(node['filename']).resolve() == expected:
                assert node.get('format') == 'raw' and node.get('virtual-size') == expected.stat().st_size, 'Unexpected source backing geometry'
                matches.append(block)
                break
            node = node.get('backing-image', {})
    assert len(matches) == 1, 'QMP does not identify exactly one source image backing device'


def run(options):
    from PIL import Image
    image = pathlib.Path(options.image).resolve(strict=True)
    receipt_path = pathlib.Path(options.boot_result).resolve(strict=True)
    receipt = json.loads(receipt_path.read_text(encoding='utf-8-sig'))
    assert receipt.get('Passed') is True and receipt.get('ImageUnchanged') is True and receipt.get('QemuExitCode') == 0, 'The image has no passing prior boot receipt'
    assert image.is_file() and image.suffix.lower() == '.img' and image.stat().st_size == 128 * 1024 * 1024, 'Only the project 128 MiB file image is accepted'
    image_hash, receipt_hash = sha256(image), sha256(receipt_path)
    assert image_hash == receipt['ImageSha256'], 'Image does not match the passing receipt'
    image_verification, files = inspect_image(image)
    assert image_verification['Files']['System/Kernel.elf'] == receipt['KernelSha256'], 'Embedded kernel does not match receipt'
    assert b'path: boot():/System/Kernel.elf' in files['Limine.conf'], 'Unexpected boot path'
    files.clear()  # Do not retain or duplicate the 100+ MiB embedded ELF.
    snapshot_path = pathlib.Path(receipt['SourceSnapshotPath']).resolve(strict=True)
    assert sha256(snapshot_path) == receipt['SourceSnapshotSha256'], 'Compiled-source snapshot hash differs'
    compiled_sources = json.loads(snapshot_path.read_text(encoding='utf-8-sig'))
    for name in ('Source/XFESystem.Kernel/NativeDesktopInput.cs', 'Source/XFESystem.Kernel/NativeKeyboardCapture.cs'):
        assert compiled_sources[name] == sha256(ROOT / name), 'Current shortcut source differs from verified compiled image: ' + name
    evidence_root = pathlib.Path(options.evidence_root).resolve() if options.evidence_root else ROOT / 'build' / 'CSharpBoot' / 'UsbShortcuts'
    parent = evidence_root
    while not parent.exists():
        parent = parent.parent
    assert shutil.disk_usage(parent).free >= 10 * 1024 ** 3, 'Need at least 10 GiB free; no test artifacts were created'
    out = evidence_root / ('Proof-' + str(time.time_ns()))
    out.mkdir(parents=True, exist_ok=False)
    firmware = ROOT / '.tools/firmware/edk2-i386-vars.fd'
    shutil.copyfile(firmware, out / 'FirmwareVars.fd')
    with socket.socket() as probe:
        probe.bind(('127.0.0.1', 0))
        port = probe.getsockname()[1]
    edid = receipt.get('EdidResolution', '3840x2160')
    assert re.fullmatch(r'[0-9]+x[0-9]+', edid), 'Malformed prior EDID resolution'
    width, height = (int(value) for value in edid.split('x'))
    assert 320 <= width <= 16384 and 200 <= height <= 16384
    memory = int(receipt.get('MemoryMiB', 1024))
    assert memory in (512, 1024, 2048)
    guest_uuid = str(uuid.uuid4())
    # A disabled i8042 leaves exactly one QEMU key sink. Thus input-send-event
    # cannot accidentally pass through the PS/2 keyboard used by older tests.
    args = [str(ROOT / '.tools/qemu/qemu-system-x86_64.exe'), '-uuid', guest_uuid, '-machine', 'q35,i8042=off,usb=off',
            '-accel', options.accel, '-cpu', 'max', '-m', str(memory) + 'M', '-smp', '8',
            '-vga', 'std', '-global', 'VGA.vgamem_mb=64', '-global', 'VGA.xres=' + str(width),
            '-global', 'VGA.yres=' + str(height), '-global', 'VGA.xmax=' + str(width),
            '-global', 'VGA.ymax=' + str(height), '-global', 'VGA.edid=on',
            '-drive', 'if=pflash,format=raw,readonly=on,file=' + str(ROOT / '.tools/firmware/edk2-x86_64-code.fd'),
            '-drive', 'if=pflash,format=raw,file=' + str(out / 'FirmwareVars.fd'),
            '-drive', 'file=' + str(image) + ',format=raw,if=ide,snapshot=on',
            '-device', 'qemu-xhci,id=XfeUsbXhci', '-device', 'usb-kbd,id=XfeUsbKeyboard,bus=XfeUsbXhci.0',
            '-nic', 'none', '-display', 'none', '-serial', 'file:' + str(out / 'Serial.log'),
            '-monitor', 'none', '-qmp', 'tcp:127.0.0.1:' + str(port) + ',server=on,wait=off', '-no-reboot']
    result = {'Kind': 'NativeUsbDesktopShortcutProof', 'Passed': False, 'PhysicalKeyboardVerified': False,
              'ImagePath': str(image), 'ImageCopied': False, 'KernelCopied': False,
              'ImageSha256': image_hash, 'PriorBootResult': str(receipt_path), 'PriorBootResultSha256': receipt_hash,
              'KernelSha256': receipt['KernelSha256'], 'SourceSnapshotSha256': receipt['SourceSnapshotSha256'],
              'HarnessSha256': sha256(__file__), 'QemuArgs': args, 'Acceleration': options.accel,
              'MemoryMiB': memory, 'AdvertisedCpuCount': 8, 'Ps2ControllerEnabled': False,
              'UsbController': 'qemu-xhci', 'KeyboardDevice': 'usb-kbd', 'Checks': []}
    result['PrivateQemuUuid'] = guest_uuid
    process, client, stream = None, None, None
    commands = []
    stdout = (out / 'QemuStdout.log').open('wb')
    stderr = (out / 'QemuStderr.log').open('wb')

    def log():
        serial = out / 'Serial.log'
        return serial.read_text(encoding='utf-8', errors='replace') if serial.exists() else ''

    def healthy():
        assert process is not None and process.poll() is None, 'Owned QEMU process exited'
        current = log()
        assert not any(marker in current for marker in FAILURES), 'Guest runtime failure; inspect Serial.log'
        serial = out / 'Serial.log'
        assert not serial.exists() or serial.stat().st_size < 8 * 1024 * 1024, 'Unexpected log growth; stopping owned guest'
        return current

    def qmp(command, arguments=None):
        identifier = len(commands) + 1
        request = {'execute': command, 'arguments': arguments or {}, 'id': identifier}
        commands.append(request)
        stream.write((json.dumps(request) + '\n').encode()); stream.flush()
        while True:
            line = stream.readline()
            assert line, 'Owned QMP connection closed'
            reply = json.loads(line)
            if 'event' in reply:
                continue
            assert reply.get('id') == identifier, 'Unexpected QMP response identity'
            assert 'error' not in reply, str(reply.get('error'))
            assert 'return' in reply, 'Missing QMP return'
            return reply['return']

    def wait(predicate, description, start=0, timeout=12):
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            current = healthy()
            value = predicate(current, start)
            if value is not None and value is not False:
                return value
            time.sleep(.04)
        raise AssertionError(description + ': no matching guest evidence\n' + log()[start:][-4000:])

    def state(predicate, description, start):
        return wait(lambda text, offset: next((item for item in reversed(snapshots(text)) if item['SerialOffset'] >= offset and predicate(item)), None), description, start)

    def key(code, down):
        qmp('input-send-event', {'events': [{'type': 'key', 'data': {'down': down, 'key': {'type': 'qcode', 'data': code}}}]})
        time.sleep(.12)

    def tap(code):
        key(code, True); key(code, False)

    def command(name, start):
        return wait(lambda text, offset: re.search(r'^XFE_NATIVE_SHORTCUT\|Desktop\|[^\r\n]*\|Key\|' + re.escape(name) + r'\|[^\r\n]*\n', text[offset:], re.M), 'Desktop shortcut ' + name, start)

    def screenshot(name):
        ppm, png = out / (name + '.ppm'), out / (name + '.png')
        qmp('screendump', {'filename': str(ppm)})
        with Image.open(ppm) as source:
            source.thumbnail((1280, 720), Image.Resampling.LANCZOS)
            source.save(png)
        ppm.unlink()
        return {'Path': str(png), 'Sha256': sha256(png), 'Meaning': 'Downscaled real guest screenshot; native geometry remains in Serial.log'}

    try:
        environment = os.environ.copy()
        environment['TMP'] = environment['TEMP'] = str(out)
        process = subprocess.Popen(args, cwd=ROOT, stdout=stdout, stderr=stderr, env=environment)
        result['OwnedQemuPid'] = process.pid
        deadline = time.monotonic() + options.boot_timeout
        while time.monotonic() < deadline:
            assert process.poll() is None, 'Owned QEMU exited before QMP connect'
            try:
                client = socket.create_connection(('127.0.0.1', port), timeout=2)
                break
            except OSError:
                time.sleep(.1)
        assert client is not None, 'Owned QMP port did not become available'
        client.settimeout(8)
        stream = client.makefile('rwb')
        greeting = json.loads(stream.readline())
        assert 'QMP' in greeting, 'No QMP greeting from owned process'
        qmp('qmp_capabilities')
        assert qmp('query-uuid')['UUID'] == guest_uuid, 'QMP is not the uniquely owned test process'
        result['PrivateQmpIdentityVerified'] = True
        result['QemuVersion'] = qmp('query-version')
        block_devices = qmp('query-block')
        require_image_backing(block_devices, image)
        result['SourceBackingVerified'] = True
        (out / 'BlockDevices.json').write_text(json.dumps(block_devices, indent=2), encoding='utf-8')
        topology = qmp('human-monitor-command', {'command-line': 'info qtree'})
        (out / 'DeviceTree.txt').write_text(topology, encoding='utf-8')
        require_usb_only_topology(topology)
        result['UsbOnlyTopologyVerified'] = True
        wait(lambda text, _: 'XFE_CSHARP_DESKTOP_READY' in text, 'Desktop ready', timeout=options.boot_timeout)
        wait(lambda text, _: '[KeyboardManager] Registered keyboard, total: 1' in text, 'USB keyboard registered')
        result['Booted'] = True
        result['Screenshots'] = [screenshot('Desktop')]
        for meta, scan in [('meta_l', 91), ('meta_r', 92)]:
            start = len(log()); tap(meta); command('StartMenu', start)
            opened = state(lambda item: item['StartMenuOpen'] == 1, meta + ' opens Start', start)
            assert re.search(r'^XFE_NATIVE_SHORTCUT\|Raw\|Scan\|' + str(scan) + r'\|Released\|0\|', log()[start:], re.M)
            assert re.search(r'^XFE_NATIVE_SHORTCUT\|Raw\|Scan\|' + str(scan) + r'\|Released\|1\|', log()[start:], re.M)
            start = len(log()); tap(meta); state(lambda item: item['StartMenuOpen'] == 0, meta + ' closes Start', start)
            result['Checks'].append({'Name': meta + ' bare', 'Passed': True, 'Opened': opened})
        windows = []
        for index, meta in enumerate(('meta_l', 'meta_r', 'meta_l')):
            start = len(log()); key(meta, True); tap('e'); key(meta, False); command('OpenFiles', start)
            opened = state(lambda item: len(item['Windows']) == index + 1 and item['StartMenuOpen'] == 0 and item['Windows'][item['ActiveId']]['Application'] == 'Files', 'Win+E opens window ' + str(index + 1), start)
            windows.append(opened['ActiveId'])
            assert len(set(windows)) == len(windows)
        result['Checks'].append({'Name': 'Win+E with both GUI keys', 'Passed': True, 'WindowIds': windows})
        first, second, third = windows
        key('alt', True)
        start = len(log()); tap('tab'); command('SwitchWindow', start)
        one = state(lambda item: item['WindowSwitching'] == 1 and item['ActiveId'] == second, 'Alt+Tab previous window', start)
        start = len(log()); tap('tab')
        two = state(lambda item: item['WindowSwitching'] == 1 and item['ActiveId'] == first, 'Repeated Alt+Tab keeps stable MRU order', start)
        key('shift', True)
        start = len(log()); tap('tab'); command('SwitchWindowBack', start)
        reverse = state(lambda item: item['WindowSwitching'] == 1 and item['ActiveId'] == second, 'Shift+Alt+Tab reverses selection', start)
        key('shift', False)
        start = len(log()); key('alt', False); command('EndWindowSwitch', start)
        committed = state(lambda item: item['WindowSwitching'] == 0 and item['ActiveId'] == second, 'Alt release commits selected window', start)
        result['Checks'].append({'Name': 'Left Alt stable forward/reverse MRU and commit', 'Passed': True,
                                 'Forward': [one, two], 'Reverse': reverse, 'Committed': committed})
        key('alt_r', True)
        start = len(log()); tap('tab'); command('SwitchWindow', start)
        right = state(lambda item: item['WindowSwitching'] == 1 and item['ActiveId'] == third, 'Right Alt+Tab selects prior MRU', start)
        start = len(log()); key('alt_r', False); command('EndWindowSwitch', start)
        state(lambda item: item['WindowSwitching'] == 0 and item['ActiveId'] == third, 'Right Alt release commits', start)
        result['Checks'].append({'Name': 'Right Alt USB modifier and release', 'Passed': True, 'Selected': right})
        start = len(log()); key('meta_l', True); tap('d'); key('meta_l', False); command('ShowDesktop', start)
        state(lambda item: item['ActiveId'] == 0 and all(window['Minimized'] == 1 for window in item['Windows'].values()), 'Win+D minimizes all windows', start)
        start = len(log()); key('meta_l', True); tap('d'); key('meta_l', False)
        state(lambda item: item['ActiveId'] != 0 and any(window['Minimized'] == 0 for window in item['Windows'].values()), 'Win+D restores desktop windows', start)
        result['Checks'].append({'Name': 'Win+D minimize and restore', 'Passed': True})
        result['Screenshots'].append(screenshot('WindowsAfterUsbShortcuts'))
        result['GuestShortcutTrace'] = [line for line in log().splitlines() if line.startswith('XFE_NATIVE_SHORTCUT|')]
        healthy()
        qmp('quit'); process.wait(timeout=15)
        assert process.returncode == 0, 'Owned QEMU did not exit cleanly'
        result['Passed'] = True
    except Exception as error:
        result['Error'] = str(error)
        if stream is not None and process is not None and process.poll() is None:
            try:
                result['ScreenshotsOnFailure'] = [screenshot('Failure')]
            except Exception as capture_error:
                result['ScreenshotFailure'] = str(capture_error)
    finally:
        if process is not None and process.poll() is None:
            try:
                if stream is not None:
                    qmp('quit')
                process.wait(timeout=10)
            except Exception:
                process.terminate()
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    process.kill(); process.wait(timeout=5)
        if stream is not None:
            stream.close()
        if client is not None:
            client.close()
        stdout.close(); stderr.close()
        result['QemuExitCode'] = process.returncode if process is not None else None
        result['OwnedProcessExited'] = process is not None and process.poll() is not None
        result['ImageUnchanged'] = sha256(image) == image_hash
        result['PriorReceiptUnchanged'] = sha256(receipt_path) == receipt_hash
        result['SourceSnapshotUnchanged'] = sha256(snapshot_path) == receipt['SourceSnapshotSha256']
        result['CurrentShortcutSourcesUnchanged'] = all(compiled_sources[name] == sha256(ROOT / name) for name in
            ('Source/XFESystem.Kernel/NativeDesktopInput.cs', 'Source/XFESystem.Kernel/NativeKeyboardCapture.cs'))
        result['Passed'] = (result['Passed'] and result['OwnedProcessExited'] and result['QemuExitCode'] == 0 and
                            result['ImageUnchanged'] and result['PriorReceiptUnchanged'] and
                            result['SourceSnapshotUnchanged'] and result['CurrentShortcutSourcesUnchanged'])
        result['ArtifactBytes'] = sum(path.stat().st_size for path in out.rglob('*') if path.is_file())
        (out / 'QmpCommands.json').write_text(json.dumps(commands, indent=2), encoding='utf-8')
        (out / 'Result.json').write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')
        print(out)
        print(json.dumps({'Passed': result['Passed'], 'Error': result.get('Error'), 'ImageUnchanged': result['ImageUnchanged'],
                          'Checks': len(result['Checks']), 'ArtifactBytes': result['ArtifactBytes']}, ensure_ascii=False))
    return 0 if result['Passed'] else 1


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--image', required=True)
    parser.add_argument('--boot-result', required=True)
    parser.add_argument('--accel', choices=['whpx', 'tcg'], default='whpx')
    parser.add_argument('--evidence-root')
    parser.add_argument('--boot-timeout', type=int, default=120)
    arguments = parser.parse_args()
    assert 15 <= arguments.boot_timeout <= 300
    raise SystemExit(run(arguments))