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
  python3 tools/run_vm.py --selinux   # boot with the host SELinux policy (permissive)

inside the guest: log in as `root` (passwordless by default, or use
--root-password). `bctl status` / `bctl shutdown poweroff` drive the
init over its control socket; `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.
--selinux flips bajia to a dynamic build (there is no static libselinux on
Fedora), bundles libselinux/libpcre2/glibc + the loader into the initramfs,
and copies the host's /etc/selinux/<type> policy (loaded permissively). This
is the first stage of bootstrapping a real policy: boot, read the `avc:
denied` lines, refine, then flip to enforcing.
"""
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, selinux: bool = False) -> Path:
    env = dict(os.environ)
    cfg = ["python3", "configure.py"]
    if selinux:
        cfg.append("--selinux")
    if static:
        print("building bajia (statically linked)...")
        env["CXX"] = env.get("CXX", "g++") + " -static"
    r = run(cfg, 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

SELINUX_CONFIG = Path("/etc/selinux/config")

def host_selinux_type() -> str:
    if not SELINUX_CONFIG.is_file():
        return "targeted"
    for line in SELINUX_CONFIG.read_text().splitlines():
        line = line.strip()
        if line.startswith("SELINUXTYPE="):
            return line.split("=", 1)[1].strip().strip('"')
    return "targeted"

def selinux_policy_files() -> list[tuple[Path, str]]:
    """Return (host_path, initramfs_relative_path) pairs for the policy payload."""
    typ = host_selinux_type()
    base = Path("/etc/selinux") / typ
    pols = sorted((base / "policy").glob("policy.*"))
    if not pols:
        sys.exit(f"--selinux: no policy under {base / 'policy'}/ "
                 "(install selinux-policy-targeted, a Fedora SELinux host is assumed)")
    # libselinux 3.x looks for the restorecon table at contexts/files/file_contexts
    # (older versions used contexts/file_contexts); bundle whichever exists.
    fc = next((p for p in (base / "contexts" / "files" / "file_contexts",
                           base / "contexts" / "file_contexts") if p.is_file()), None)
    if not fc:
        sys.exit(f"--selinux: missing file_contexts under {base / 'contexts'}/")
    fc_rel = "etc/selinux/" + typ + "/" + str(fc.relative_to(base))
    prefix = f"etc/selinux/{typ}/"
    return [(pols[-1], prefix + f"policy/{pols[-1].name}"),
            (fc, fc_rel)]

def bundle_dynamic_libs(init: Path, root: Path) -> None:
    """Copy the dynamic loader + resolved .so deps into the initramfs,
    mirroring their absolute paths so the interpreter finds them."""
    out = run(["ldd", str(init)], capture_output=True, text=True)
    if out.returncode != 0:
        sys.exit("ldd failed on " + str(init))
    libs: list[str] = []
    for line in out.stdout.splitlines():
        line = line.strip()
        if "=>" in line:
            path = line.split("=>", 1)[1].strip().split(" ", 1)[0].strip()
        else:  # "linux-vdso" or the loader line "/lib64/ld-linux-x86-64.so.2 (0x...)"
            path = line.split(" ", 1)[0].strip()
        if path.startswith("/") and Path(path).is_file():
            libs.append(path)
    seen: set[str] = set()
    for lib in libs:  # preserve first-seen order
        if lib in seen:
            continue
        seen.add(lib)
        dst = root / lib.lstrip("/")
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(lib, dst)
        dst.chmod(0o755)
    print("bundled dynamic libs:", ", ".join(seen))
    if not seen:
        sys.exit("ldd reported no libraries - unexpected for a dynamic binary")

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"

# appended to the test init.rc; started on boot, prints the exec context of a
# seclabel'd service and of init itself, then exits.
SELINUX_RC_PROBE = """
service selinux-probe /bin/sh -c "echo probe-ctx=$(cat /proc/self/attr/current) init-ctx=$(cat /proc/1/attr/current)"
    console
    seclabel = system_u:system_r:init_t:s0
    respawn = never

on boot
    start selinux-probe
"""

def build_initramfs(init: Path, busybox: Path, rc_text: str, root_password: str | None,
                    selinux: bool, keep: bool,
                    bundles: list[tuple[str, Path]] | None = None) -> 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)

        if selinux:
            for src, rel in selinux_policy_files():
                dst = root / rel
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy(src, dst)
            selcfg = root / "etc" / "selinux" / "config"
            selcfg.write_text(f"SELINUX=permissive\nSELINUXTYPE={host_selinux_type()}\n")
            bundle_dynamic_libs(init, root)
        elif 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.")

        if selinux:
            rc_text = rc_text + SELINUX_RC_PROBE

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

        ctl = BUILD / "bctl"
        if ctl.is_file():
            shutil.copy(ctl, root / "bin" / "bctl")
            (root / "bin" / "bctl").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) +
            "nobody:x:65534:65534:nobody:/:/bin/sh\n")
        (root / "etc" / "group").write_text(
            "root:x:0:\n"
            "nobody:x:65534:\n"
            "daemon:x:1:\n")

        for rel, src in (bundles or []):
            dst = root / rel.lstrip("/")
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(src, dst)
            dst.chmod(0o644)

        # The staging tree is owned by the host user and mkdtemp makes the
        # top dir 0700; GNU cpio preserves both, so without this the guest's
        # "/" would be mode 0700 owned by uid 1000 -- fine for root services,
        # but dropped-privilege services couldn't traverse it. Stamp owner
        # root:root (no host chown needed) and make the root traversable.
        p = run(["bash", "-c",
                 "cd \"$1\" && chmod 0755 . && "
                 "find . -print0 | cpio --null -o -H newc --owner=0:0",
                 "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"
    append = (f"console=tty1 console=ttyS0 rdinit=/init loglevel={args.loglevel}")
    if args.selinux:
        append += " selinux=1 enforcing=0"
    cmd = [
        qemu,
        "-M", args.machine,
        "-m", str(args.mem),
        "-smp", str(args.smp),
        "-kernel", str(kernel),
        "-initrd", str(initrd),
        "-append", append,
        "-display", display,
        "-serial", "stdio",
    ]
    if args.nographic:
        cmd[cmd.index("-display") + 1] = "none"
    if args.serial_log:
        cmd[cmd.index("-display") + 1] = "none"
        cmd[cmd.index("-serial") + 1] = f"file:{args.serial_log}"
    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("--bundle", action="append", default=[],
                    metavar="REL=HOSTPATH",
                    help="copy HOSTPATH into the initramfs at absolute REL "
                         "(repeatable; e.g. --bundle "
                         "etc/bajia/extra.rc=/tmp/extra.rc for @import tests)")
    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")
    ap.add_argument("--selinux", action="store_true",
                    help="bundle the host SELinux policy + libs, boot permissive")
    ap.add_argument("--serial-log", type=Path,
                    help="write the serial console to this file (forces -display none)")
    args = ap.parse_args()

    static = not args.no_static and not args.selinux
    init = BAJIA if args.no_build else build_bajia(static, args.selinux)
    if not init.is_file():
        sys.exit(f"bajia not built at {init} (drop --no-build)")
    if not static and not args.selinux and 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
    bundles: list[tuple[str, Path]] = []
    for spec in args.bundle:
        rel, _, path = spec.partition("=")
        src = Path(path)
        if not src.is_file():
            sys.exit(f"--bundle: host file not found: {src}")
        if not rel.startswith("/"):
            sys.exit(f"--bundle: REL must be absolute, got: {rel!r}")
        if ".." in [c for c in Path(rel).parts]:
            sys.exit(f"--bundle: REL must not contain '..': {rel}")
        bundles.append((rel, src))
    initrd = build_initramfs(init, busybox, rc_text, args.root_password,
                             selinux=args.selinux, keep=args.keep_initramfs,
                             bundles=bundles)
    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)