vmix.nix/lib/images/macos/helpers/vm-driver.py

#!/usr/bin/env python3
"""vmix macOS VM driver: runs QEMU and decides when a build boot is finished.

Modes
  pe       the recovery PE runs /Volumes/VMIX/run.sh and powers off. Success is
           QEMU exiting on its own; the caller checks vmix-run.status.
  install  the PE starts the macOS installer, which reboots through its phases
           into the installed system. Finished when that system reaches the
           (bright) loginwindow / Setup Assistant — the driver powers it down —
           or halts on its own.
  boot     boot an installed image and wait for it to halt or reach the loginwindow.

Observation is passive: the serial console (boot-args serial=3: the PE's
"VMIX-*" markers, kernel boots, panics) and screenshots over QMP (mean
brightness + a coarse change fingerprint). No OCR, no keystrokes. A boot hang
(dark, frozen screen, disk and serial idle) is retried with a system_reset.
Screenshots, the driver log and the serial log are kept in --debug-dir.
"""
import argparse
import hashlib
import io
import json
import os
import shutil
import socket
import subprocess
import sys
import time

try:
    from PIL import Image
except ImportError:  # screenshots then only serve as debug files
    Image = None

PANIC_MARKS = ('panic(cpu', 'Kernel Extensions in backtrace', 'Debugger called: <panic>', 'Nested panic detected', 'panic string:')
REBOOT_MARK = 'MACH Reboot'


class Log:
    def __init__(self, path):
        self.f = open(path, 'a') if path else None
        self.t0 = time.time()

    def __call__(self, msg):
        line = f'[{time.time() - self.t0:7.1f}s] {msg}'
        print(f'vmix driver: {line}', flush=True)
        if self.f:
            self.f.write(line + '\n')
            self.f.flush()


class QMP:
    def __init__(self, path):
        self.path = path
        self.s = None
        self.buf = b''

    def connect(self, timeout=90):
        t0 = time.time()
        while True:
            try:
                s = socket.socket(socket.AF_UNIX)
                s.settimeout(60)
                s.connect(self.path)
                self.s = s
                self.buf = b''
                self._read()            # greeting
                self.cmd('qmp_capabilities')
                return
            except (OSError, ValueError):
                if time.time() - t0 > timeout:
                    raise
                time.sleep(1)

    def _read(self):
        while b'\n' not in self.buf:
            d = self.s.recv(65536)
            if not d:
                raise OSError('QMP socket closed')
            self.buf += d
        line, self.buf = self.buf.split(b'\n', 1)
        return json.loads(line)

    def cmd(self, name, **args):
        self.s.sendall(json.dumps({'execute': name, 'arguments': args}).encode() + b'\n')
        while True:
            r = self._read()
            if 'return' in r:
                return r['return']
            if 'error' in r:
                raise RuntimeError(r['error'])

    def screendump(self, path):
        self.cmd('screendump', filename=path)

    def system_reset(self):
        self.cmd('system_reset')

    def system_powerdown(self):
        self.cmd('system_powerdown')


class Screen:
    """Screenshots over QMP with a coarse change fingerprint (cursor-insensitive)."""

    def __init__(self, qmp, debug_dir, log):
        self.qmp = qmp
        self.dir = debug_dir
        self.log = log
        self.tmp = os.path.join(debug_dir, f'.grab-{os.getpid()}.ppm')
        self.img = None
        self.last_fp = None
        self.stable_since = time.time()
        self.n = 0

    def grab(self):
        self.qmp.screendump(self.tmp)
        with open(self.tmp, 'rb') as f:
            data = f.read()
        if Image is None:
            fp = hashlib.sha256(data).hexdigest()
        else:
            self.img = Image.open(io.BytesIO(data))
            small = self.img.convert('L').resize((48, 36))
            fp = hashlib.sha256(bytes(b & 0xF0 for b in small.tobytes())).hexdigest()
        if fp != self.last_fp:
            self.last_fp = fp
            self.stable_since = time.time()
        return self.img

    def stable_for(self):
        return time.time() - self.stable_since

    def mean(self):
        if self.img is None:
            return 0
        g = self.img.convert('L').resize((64, 48))
        px = g.tobytes()
        return sum(px) / len(px)

    def is_blank(self):
        return self.img is not None and self.mean() < 3

    def save(self, tag):
        self.n += 1
        path = os.path.join(self.dir, f'{self.n:03d}-{tag}.png')
        try:
            if self.img is not None:
                self.img.save(path)
            else:
                shutil.copy(self.tmp, path.replace('.png', '.ppm'))
        except Exception as e:  # noqa: BLE001
            self.log(f'could not save screenshot: {e}')
        return path


class Serial:
    """Tail the serial console file QEMU writes (-serial file:...)."""

    def __init__(self, path):
        self.path = path
        self.pos = 0
        self.last_activity = time.time()
        self.boots = 0
        self.reboot_at = None

    def poll(self):
        if not self.path or not os.path.exists(self.path):
            return []
        with open(self.path, 'rb') as f:
            f.seek(self.pos)
            data = f.read()
            self.pos = f.tell()
        if not data:
            return []
        self.last_activity = time.time()
        lines = data.decode('utf-8', 'replace').replace('\r', '').split('\n')
        for l in lines:
            if l.startswith('Darwin Kernel Version'):
                self.boots += 1
                self.reboot_at = None
            elif REBOOT_MARK in l:
                self.reboot_at = time.time()
        return lines

    def idle_for(self):
        return time.time() - self.last_activity


def prepare_debug_dir(path):
    os.makedirs(path, exist_ok=True)
    try:
        os.chmod(path, 0o777)
    except OSError:
        pass
    for f in os.listdir(path):
        if f.endswith(('.png', '.ppm', '.log')) or f.startswith('.grab-') or f == 'qmp.sock':
            try:
                os.remove(os.path.join(path, f))
            except OSError:
                pass


def launch(qemu_args, qmp_sock, log):
    if os.path.exists(qmp_sock):
        os.remove(qmp_sock)
    cmd = list(qemu_args) + ['-qmp', f'unix:{qmp_sock},server,nowait']
    log('launching: ' + ' '.join(cmd))
    return subprocess.Popen(cmd)


def disk_idle(args):
    if not args.progress_file:
        return True
    try:
        return (time.time() - os.path.getmtime(args.progress_file)) > args.disk_idle
    except OSError:
        return True


def drive(args, proc, qmp, screen, serial, log):
    start = time.time()
    last_periodic = 0
    resets = 0
    panics = 0
    started = args.mode == 'boot'      # pe/install: wait for the PE marker first
    blank_since = None
    login_since = None
    while True:
        rc = proc.poll()
        if rc is not None:
            log(f'QEMU exited with {rc}')
            return rc
        now = time.time()
        if now - start > args.timeout:
            try:
                screen.grab()
                screen.save('timeout')
            except Exception:  # noqa: BLE001
                pass
            log('timeout reached, killing QEMU')
            proc.kill()
            return 124
        time.sleep(args.interval)

        panic = False
        for line in serial.poll():
            if 'VMIX' in line:
                log('serial: ' + line.strip()[:220])
                if 'VMIX-PE: running run.sh' in line and not started:
                    started = True
                    log('PE started run.sh')
                if 'VMIX-PE: no VMIX volume' in line and args.mode != 'boot':
                    log('PE did not find the VMIX volume')
                    proc.kill()
                    return 3
            if line.startswith('Darwin Kernel Version'):
                log(f'guest kernel boot #{serial.boots}')
            if any(m in line for m in PANIC_MARKS):
                panic = True
                log('serial: ' + line.strip()[:220])
        if panic:
            panics += 1
            try:
                screen.grab()
                screen.save('panic')
            except Exception:  # noqa: BLE001
                pass
            if args.mode == 'pe' or panics > args.max_resets:
                log(f'kernel panic #{panics}, giving up')
                proc.kill()
                return 3
            log(f'kernel panic #{panics}, system_reset')
            qmp.system_reset()
            screen.stable_since = time.time()
            continue

        # The guest asked for a reboot but no kernel came back: macOS' restart
        # path panics in QEMU (AppleSMC watchdog keys, see README); reset now
        # instead of waiting for the frozen-screen watchdog.
        if serial.reboot_at and now - serial.reboot_at > args.reboot_timeout and args.mode != 'pe':
            resets += 1
            try:
                screen.grab()
                screen.save('reboot-dead')
            except Exception:  # noqa: BLE001
                pass
            log(f'guest requested a reboot {now - serial.reboot_at:.0f}s ago and died, system_reset #{resets}')
            serial.reboot_at = None
            if resets > args.max_resets:
                proc.kill()
                return 3
            qmp.system_reset()
            screen.stable_since = time.time()
            continue

        if not started and now - start > args.start_timeout:
            try:
                screen.grab()
                screen.save('no-start')
            except Exception:  # noqa: BLE001
                pass
            log(f'PE did not start run.sh within {args.start_timeout:.0f}s')
            proc.kill()
            return 3

        try:
            screen.grab()
        except Exception as e:  # noqa: BLE001
            log(f'screendump failed ({e})')
            time.sleep(2)
            continue
        if now - last_periodic > args.periodic:
            last_periodic = now
            screen.save('periodic')

        if args.mode == 'pe':
            continue                    # the PE powers off by itself; nothing to decide
        # install: the PE phase (kernel boot #1) is protected by the guest's own
        # retries; the checks below apply once the installer has rebooted.
        in_os = args.mode == 'boot' or serial.boots >= 2
        if not in_os or screen.stable_for() < args.settle:
            continue
        idle = disk_idle(args) and serial.idle_for() > args.disk_idle

        if screen.is_blank():
            blank_since = blank_since or now
            # macOS `shutdown -h` halts to a black screen without an ACPI power-off
            if now - blank_since > args.halt_timeout and idle:
                screen.save('halt')
                log(f'guest halted (black {now - blank_since:.0f}s, idle); killing QEMU')
                proc.kill()
                return 0
            continue
        blank_since = None

        if screen.mean() > args.bright:
            # loginwindow / Setup Assistant: the OS is installed and booted
            if login_since is None:
                login_since = now
                log(f'bright screen (mean {screen.mean():.0f}): loginwindow/desktop, grace {args.login_grace:.0f}s')
            elif now - login_since > args.login_grace:
                screen.save('loginwindow')
                log('powering down (boot complete)')
                try:
                    qmp.system_powerdown()
                except Exception as e:  # noqa: BLE001
                    log(f'powerdown failed: {e}')
                for _ in range(120):
                    if proc.poll() is not None:
                        log('QEMU exited after powerdown')
                        return 0
                    time.sleep(1)
                log('guest ignored powerdown, killing QEMU')
                proc.kill()
                return 0
            continue
        login_since = None

        # dark, frozen, nothing happening: a boot hang (seen at the Apple logo)
        if screen.stable_for() > args.stall_reset and idle and resets < args.max_resets:
            resets += 1
            screen.save('stall-reset')
            log(f'boot hung ({screen.stable_for():.0f}s frozen, dark, idle), system_reset #{resets}')
            try:
                qmp.system_reset()
            except Exception as e:  # noqa: BLE001
                log(f'system_reset failed: {e}')
            screen.stable_since = time.time()


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--mode', choices=['pe', 'install', 'boot'], required=True)
    p.add_argument('--name', default='macos')
    p.add_argument('--debug-dir', default=None)
    p.add_argument('--serial-log', default=None, help='file QEMU writes the serial console to')
    p.add_argument('--timeout', type=int, default=4 * 3600, help='seconds before QEMU is killed')
    p.add_argument('--start-timeout', type=float, default=600.0, help='seconds for the PE to start run.sh')
    p.add_argument('--interval', type=float, default=5.0, help='seconds between screenshots')
    p.add_argument('--periodic', type=float, default=120.0, help='seconds between saved debug screenshots')
    p.add_argument('--settle', type=float, default=12.0, help='seconds a screen must be unchanged to act on it')
    p.add_argument('--stall-reset', type=float, default=360.0, help='reset the VM if a dark screen is frozen this long while disk and serial are idle')
    p.add_argument('--max-resets', type=int, default=6)
    p.add_argument('--reboot-timeout', type=float, default=60.0, help='seconds after a guest reboot request without a new kernel boot before the VM is reset')
    p.add_argument('--halt-timeout', type=float, default=150.0, help='a pure-black, idle screen this long means the guest halted')
    p.add_argument('--progress-file', default=None, help='the system disk; its mtime shows guest disk activity')
    p.add_argument('--disk-idle', type=float, default=90.0, help='seconds without disk/serial activity that count as idle')
    p.add_argument('--bright', type=float, default=80.0, help='mean brightness above which a screen is the loginwindow/desktop')
    p.add_argument('--login-grace', type=float, default=180.0, help='seconds a bright screen must persist before powering down')
    p.add_argument('qemu', nargs=argparse.REMAINDER)
    args = p.parse_args()
    qemu_args = [a for a in args.qemu if a != '--']
    if not qemu_args:
        p.error('QEMU command line required after --')

    debug_dir = args.debug_dir or f'/tmp/vmix-macos/{args.name}'
    prepare_debug_dir(debug_dir)
    log = Log(os.path.join(debug_dir, 'driver.log'))
    log(f'mode={args.mode} debug-dir={debug_dir} serial={args.serial_log}')

    qmp_sock = os.path.join(debug_dir, 'qmp.sock')
    proc = launch(qemu_args, qmp_sock, log)
    qmp = QMP(qmp_sock)
    try:
        qmp.connect()
    except Exception as e:  # noqa: BLE001
        log(f'QMP connect failed: {e}')
        proc.kill()
        return 2
    screen = Screen(qmp, debug_dir, log)
    serial = Serial(args.serial_log)
    try:
        rc = drive(args, proc, qmp, screen, serial, log)
    finally:
        if args.serial_log and os.path.exists(args.serial_log):
            try:
                shutil.copy(args.serial_log, os.path.join(debug_dir, 'serial.log'))
            except OSError:
                pass
        try:
            if proc.poll() is None:
                proc.kill()
        except Exception:  # noqa: BLE001
            pass
    log(f'done rc={rc}')
    return rc


if __name__ == '__main__':
    sys.exit(main())