bajia/tools/run_vm.py

#!/usr/bin/env python3
"""boot bajia as PID 1 in a QEMU VM using a tiny initramfs.

sets up everything needed for a realistic test run:
  - builds bajia (optionally statically linked)
  - builds an initramfs with bajia as /init, busybox, and a test init.rc
  - runs qemu-system-x86_64 with a GTK display, serial console, and a gdb stub

examples:
  python3 tools/run_vm.py --busybox /path/to/busybox-static
  python3 tools/run_vm.py --fetch-busybox --nographic
  python3 tools/run_vm.py --kernel /boot/vmlinuz-$(uname -r) --gdb
  python3 tools/run_vm.py --config my-init.rc

inside the guest: log in as `root` (passwordless by default, or use
--root-password), then `kill -TERM 1` -> reboot path, `kill -INT 1` ->
poweroff path.
bajia is built statically by default: a dynamic binary cannot exec inside the
initramfs (no libc there). Use --no-static only if you ship the libs too.
"""
from __future__ import annotations

import argparse
import glob
import gzip
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
BUILD = ROOT / "build"
BAJIA = BUILD / "bajia"

DEFAULT_RC = """\
# generated by tools/run_vm.py - minimal init.rc for VM testing.
on early-init
    mount proc /proc proc
    mount sysfs /sys sysfs
    mount devtmpfs /dev devtmpfs
    mkdir /dev/pts 0755
    mount devpts /dev/pts devpts
    mkdir /run 0755
    mount tmpfs /run tmpfs

on init
    write /proc/sys/kernel/hostname bajia
    log **** bajia init on-line ****

on boot
    start console-serial
    start console-tty1

on shutdown
    log shutdown: stopping services

service console-serial /bin/getty -L ttyS0 115200 vt100
    console
    respawn always

service console-tty1 /bin/getty -L 38400 tty1 vt100
    console
    respawn always
"""

# source tarballs of busybox (github.com/mirror/busybox); a pinned tag is
# tried first, then the master branch. built statically into $XDG_CACHE_HOME.
BUSYBOX_URLS = [
    "https://github.com/mirror/busybox/archive/refs/tags/1_36_1.tar.gz",
    "https://github.com/mirror/busybox/archive/refs/heads/master.tar.gz",
]

BUSYBOX_APPLETS = ["sh", "getty", "mount", "sync", "ls", "cat", "kill", "ps",
                   "poweroff", "reboot", "mkdir", "mknod", "login"]


def run(cmd, **kw) -> subprocess.CompletedProcess:
    print("$", " ".join(str(c) for c in cmd))
    return subprocess.run(cmd, **kw)


def build_bajia(static: bool) -> Path:
    env = dict(os.environ)
    if static:
        print("building bajia (statically linked)...")
        env["CXX"] = env.get("CXX", "g++") + " -static"
    r = run(["python3", "configure.py"], env=env)
    if r.returncode != 0:
        sys.exit("configure.py failed")
    r = run(["ninja", "-C", str(BUILD)])
    if r.returncode != 0:
        sys.exit("ninja build failed")
    return BAJIA

def find_kernel() -> Path | None:
    p = Path("/boot/vmlinuz-" + os.uname().release)
    if p.is_file():
        return p
    matches = sorted(glob.glob("/boot/vmlinuz-*"))
    return Path(matches[-1]) if matches else None

def is_static(path: Path) -> bool:
    filetool = shutil.which("file")
    if not filetool:
        return True  # cannot tell; don't nag
    out = run([filetool, "-b", str(path)], capture_output=True, text=True)
    blob = out.stdout.lower()
    return "static" in blob or "statically" in blob

def find_busybox() -> Path | None:
    for cand in (os.environ.get("BAJIA_BUSYBOX"), shutil.which("busybox")):
        if cand and Path(cand).is_file():
            return Path(cand)
    return None

def fetch_busybox(cache: Path) -> Path:
    """download busybox source from github.com/mirror/busybox and build it
    statically. The resulting binary is cached in `cache`."""
    dst = cache / "busybox"
    if dst.is_file():
        print("using cached busybox:", dst)
        return dst

    src_dir = cache / "busybox-src"
    src_dir.mkdir(parents=True, exist_ok=True)
    tree = None
    for url in BUSYBOX_URLS:
        try:
            print("downloading busybox source:", url)
            req = urllib.request.Request(url, headers={"User-Agent": "bajia-run-vm"})
            with urllib.request.urlopen(req, timeout=120) as resp, open(
                    cache / "busybox.tar.gz", "wb") as out:
                shutil.copyfileobj(resp, out)
            print("extracting busybox source...")
            with tarfile.open(cache / "busybox.tar.gz", "r:gz") as tf:
                tf.extractall(src_dir)
            tops = [p for p in src_dir.iterdir() if p.is_dir()]
            if len(tops) == 1:
                tree = tops[0]
                break
            print("  unexpected source layout, trying next URL")
        except Exception as e:  # noqa: BLE001 - try the next URL
            print("  failed:", e)
    if not tree:
        sys.exit("could not fetch busybox source from github.com/mirror/busybox")

    env = dict(os.environ)
    r = run(["make", "defconfig"], cwd=tree, env=env)
    if r.returncode != 0:
        sys.exit("busybox defconfig failed")
    r = run(["sed", "-i", "s|# CONFIG_STATIC is not set|CONFIG_STATIC=y|",
             tree / ".config"])
    if r.returncode != 0:
        sys.exit("busybox config edit failed")
    # CONFIG_TC needs the CBQ uapi structs (tc_cbq_*) that modern kernel
    # headers no longer provide; not needed in an initramfs, so drop it.
    r = run(["sed", "-i", "s|^CONFIG_TC=y$|# CONFIG_TC is not set|",
             tree / ".config"])
    if r.returncode != 0:
        sys.exit("busybox config edit failed")
    r = run(["make", f"-j{os.cpu_count() or 2}", "busybox"], cwd=tree, env=env)
    if r.returncode != 0:
        sys.exit("busybox build failed; if another applet breaks on your "
                 "kernel headers, disable it in busybox-src/.config and pass "
                 "--busybox with a prebuilt static binary instead")

    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy(tree / "busybox", dst)
    dst.chmod(0o755)
    print("built busybox:", dst)
    return dst

def default_cache_dir() -> Path:
    base = Path(os.environ.get("XDG_CACHE_HOME",
                               str(Path.home() / ".cache")))
    return base / "bajia"

def crypt_password(password: str) -> str:
    try:
        import crypt
        return crypt.crypt(password, crypt.mksalt(crypt.METHOD_SHA512))
    except (ImportError, AttributeError):
        p = subprocess.run(["openssl", "passwd", "-6", password],
                           capture_output=True, text=True)
        if p.returncode != 0:
            sys.exit("cannot hash the root password: need python `crypt` or openssl")
        return p.stdout.strip()

def root_passwd_line(password: str | None) -> str:
    field = crypt_password(password) if password else ""
    return f"root:{field}:0:0:root:/:/bin/sh\n"

def build_initramfs(init: Path, busybox: Path, rc_text: str, root_password: str | None,
                    keep: bool) -> Path:
    if not shutil.which("cpio"):
        sys.exit("cpio not found (install cpio)")
    root = Path(tempfile.mkdtemp(prefix="bajia-root-"))
    try:
        for sub in ("etc/bajia", "bin", "sbin", "usr/sbin", "usr/bin",
                    "dev", "proc", "sys", "run", "tmp"):
            (root / sub).mkdir(parents=True)

        shutil.copy(init, root / "init")
        (root / "init").chmod(0o755)

        shutil.copy(busybox, root / "bin" / "busybox")
        (root / "bin" / "busybox").chmod(0o755)
        for applet in BUSYBOX_APPLETS:
            (root / "bin" / applet).symlink_to("busybox")

        (root / "etc" / "bajia" / "init.rc").write_text(rc_text)
        (root / "etc" / "passwd").write_text(root_passwd_line(root_password))

        p = run(["bash", "-c",
                 "cd \"$1\" && find . -print0 | cpio --null -o -H newc",
                 "bajia-initramfs", str(root)], stdout=subprocess.PIPE)
        if p.returncode != 0:
            sys.exit("cpio packing failed")
        initrd = Path(tempfile.gettempdir()) / "bajia-initrd.cpio.gz"
        initrd.write_bytes(gzip.compress(p.stdout))

        if keep:
            print("initramfs root tree kept at:", root)
        return initrd
    finally:
        if not keep:
            shutil.rmtree(root, ignore_errors=True)

def qemu_command(kernel: Path, initrd: Path, args: argparse.Namespace) -> list[str]:
    qemu = args.qemu or shutil.which("qemu-system-x86_64") or "qemu-system-x86_64"
    display = args.display
    if display is None:
        display = "gtk" if os.environ.get("DISPLAY") else "none"
    cmd = [
        qemu,
        "-M", args.machine,
        "-m", str(args.mem),
        "-smp", str(args.smp),
        "-kernel", str(kernel),
        "-initrd", str(initrd),
        "-append",
        f"console=tty1 console=ttyS0 rdinit=/init loglevel={args.loglevel}",
        "-display", display,
        "-serial", "stdio",
    ]
    if args.nographic:
        cmd[cmd.index("-display") + 1] = "none"
    if args.gdb or args.wait_gdb:
        cmd += ["-gdb", "tcp::1234", "-S"] if args.wait_gdb else ["-s"]
    cmd += ["-no-reboot", "-no-shutdown"]
    return cmd

def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--kernel", help="kernel bzImage (default: host /boot/vmlinuz-*)")
    ap.add_argument("--busybox", help="static busybox binary to bundle")
    ap.add_argument("--fetch-busybox", action="store_true",
                    help="download busybox source (github.com/mirror/busybox), "
                         "build it static, cache in ~/.cache/bajia")
    ap.add_argument("--config", type=Path,
                    help="use this init.rc instead of the bundled test config")
    ap.add_argument("--root-password", default=None,
                    help="password for the root account in the guest "
                         "(default: passwordless login)")
    ap.add_argument("--no-static", action="store_true",
                    help="build bajia dynamically linked (won't exec in the "
                         "initramfs unless you also pack the libs)")
    ap.add_argument("--no-build", action="store_true",
                    help="use the existing build/bajia without rebuilding")
    ap.add_argument("--display", choices=["gtk", "sdl", "none"],
                    help="QEMU display backend (default: gtk if $DISPLAY set)")
    ap.add_argument("--nographic", action="store_true",
                    help="headless; serial console on stdio")
    ap.add_argument("--machine", default="q35", help="QEMU machine type")
    ap.add_argument("--mem", type=int, default=512, help="RAM in MB")
    ap.add_argument("--smp", type=int, default=2, help="virtual CPUs")
    ap.add_argument("--loglevel", type=int, default=4,
                    help="kernel loglevel (4=dmesg, 7=everything)")
    ap.add_argument("--qemu", help="qemu binary (default: qemu-system-x86_64)")
    ap.add_argument("--gdb", action="store_true", help="expose a gdb stub on :1234")
    ap.add_argument("--wait-gdb", action="store_true",
                    help="pause the machine until a gdb client attaches")
    ap.add_argument("--keep-initramfs", action="store_true",
                    help="don't delete the initramfs staging tree")
    args = ap.parse_args()

    init = BAJIA if args.no_build else build_bajia(not args.no_static)
    if not init.is_file():
        sys.exit(f"bajia not built at {init} (drop --no-build)")
    if not is_static(init):
        print("warning: bajia is dynamically linked; /init will fail to exec "
              "inside the initramfs (error -2). Rebuild with --no-static "
              "unset (static is the default) or drop --no-build.")

    kernel = Path(args.kernel) if args.kernel else find_kernel()
    if not kernel or not kernel.is_file():
        sys.exit("no kernel found: pass --kernel /path/to/bzImage or install a "
                 "kernel and use its /boot/vmlinuz-*")

    if args.fetch_busybox:
        busybox = fetch_busybox(default_cache_dir())
    else:
        busybox = Path(args.busybox) if args.busybox else find_busybox()
    if not busybox or not busybox.is_file():
        sys.exit("no busybox found: pass --busybox, --fetch-busybox, or install "
                 "busybox-static (Debian/Ubuntu: `apt install busybox-static`)")
    if not is_static(busybox):
        print("warning: busybox looks dynamically linked; a static build is "
              "safest inside an initramfs")

    rc_text = args.config.read_text() if args.config else DEFAULT_RC
    initrd = build_initramfs(init, busybox, rc_text, args.root_password,
                             keep=args.keep_initramfs)
    print("initramfs:", initrd, f"({initrd.stat().st_size / 1024:.0f} KB)")

    cmd = qemu_command(kernel, initrd, args)
    print("$", " ".join(cmd))
    return subprocess.run(cmd).returncode

if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)