cloudhost/run-ubuntu2604.py

#!/usr/bin/env python3
"""Start the Ubuntu 26.04 virtual machine."""

import os
import platform
import re
import shutil
import string
import subprocess
import sys
from pathlib import Path

script_dir = Path(__file__).resolve().parent
env = os.environ.get


def die(message, code=1):
    print(f"Error: {message}", file=sys.stderr)
    sys.exit(code)


def run(command):
    """Run a command, exiting with its status on failure (like `set -e`)."""
    result = subprocess.run(command)
    if result.returncode != 0:
        sys.exit(result.returncode)


def is_executable(path):
    return path.is_file() and os.access(path, os.X_OK)


if env("VM_IMAGE_HELPER"):
    image_downloader = Path(env("VM_IMAGE_HELPER"))
elif is_executable(script_dir / "vm-images" / "download-images.sh"):
    image_downloader = script_dir / "vm-images" / "download-images.sh"
else:
    image_downloader = script_dir / ".." / "vm-images" / "download-images.sh"

uses_shared_base_image = False
if env("UBUNTU_VM_BASE_IMAGE"):
    base_image = Path(env("UBUNTU_VM_BASE_IMAGE"))
else:
    if not is_executable(image_downloader):
        die(f"shared image helper is unavailable: {image_downloader}")
    result = subprocess.run(
        [str(image_downloader), "--print-path", "ubuntu2604"],
        stdout=subprocess.PIPE,
        text=True,
    )
    if result.returncode != 0:
        sys.exit(result.returncode)
    base_image = Path(result.stdout.rstrip("\n"))
    uses_shared_base_image = True

state_dir = Path(env("UBUNTU_VM_STATE_DIR") or script_dir / ".vm" / "ubuntu2604")
disk_size = env("UBUNTU_VM_DISK_SIZE") or "20G"


def normalize_arch(arch):
    if arch in ("arm64", "aarch64"):
        return "arm64"
    if arch in ("x86_64", "amd64"):
        return "amd64"
    die(f"unsupported architecture: {arch}")


native_arch = normalize_arch(platform.machine())
host_arch = normalize_arch(env("VM_IMAGE_ARCH") or native_arch)
qemu_binary = "qemu-system-aarch64" if host_arch == "arm64" else "qemu-system-x86_64"
qemu_config_template = Path(
    env("UBUNTU_VM_QEMU_CONFIG") or script_dir / "qemu" / "ubuntu2604.ini")
host_os = platform.system().lower()
if host_os == "darwin" and host_arch == "amd64":
    die("amd64 guests are not supported on macOS hosts; use an arm64 image")
qemu_platform_config = script_dir / "qemu" / f"ubuntu2604-{host_os}-{host_arch}.ini"


def pick_firmware(override, *candidates):
    if override:
        return Path(override)
    for candidate in candidates:
        if Path(candidate).is_file():
            return Path(candidate)
    return Path(candidates[0])


if host_arch == "arm64":
    uefi_code = pick_firmware(
        env("UBUNTU_VM_UEFI_CODE"),
        "/usr/share/AAVMF/AAVMF_CODE.fd",
        "/usr/share/edk2/aarch64/QEMU_EFI.fd",
        "/opt/homebrew/share/qemu/edk2-aarch64-code.fd",
        "/usr/local/share/qemu/edk2-aarch64-code.fd",
    )
    uefi_vars_template = pick_firmware(
        env("UBUNTU_VM_UEFI_VARS"),
        "/usr/share/AAVMF/AAVMF_VARS.fd",
        "/usr/share/edk2/aarch64/vars-template-pflash.raw",
        "/opt/homebrew/share/qemu/edk2-arm-vars.fd",
        "/usr/local/share/qemu/edk2-arm-vars.fd",
    )
else:
    uefi_code = pick_firmware(
        env("UBUNTU_VM_UEFI_CODE"),
        "/usr/share/OVMF/OVMF_CODE_4M.fd",
        "/usr/share/edk2/ovmf/OVMF_CODE.fd",
        "/opt/homebrew/share/qemu/edk2-x86_64-code.fd",
        "/usr/local/share/qemu/edk2-x86_64-code.fd",
    )
    uefi_vars_template = pick_firmware(
        env("UBUNTU_VM_UEFI_VARS"),
        "/usr/share/OVMF/OVMF_VARS_4M.fd",
        "/usr/share/edk2/ovmf/OVMF_VARS.fd",
        "/opt/homebrew/share/qemu/edk2-i386-vars.fd",
        "/usr/local/share/qemu/edk2-i386-vars.fd",
    )
daemonize = False
download_only = False
extra_args = []


def usage(stream=sys.stdout):
    stream.write(f"""\
Usage: {Path(sys.argv[0]).name} [--daemon] [--download-only] [--state-dir DIR] [-- QEMU_ARGUMENT ...]

Start an Ubuntu 26.04 LTS cloud image virtual machine. A writable overlay is
created in the state directory (default: .vm/ubuntu2604/), so the downloaded
base image is never modified. The VM hardware (memory, CPUs, SSH port forward,
...) is defined in qemu/ubuntu2604.ini, which is rendered to qemu.ini in the
state directory and loaded with QEMU's -readconfig, followed by
qemu/ubuntu2604-<os>-<arch>.ini (os is linux or darwin, arch is amd64 or arm64). Later QEMU arguments override them, e.g.
`-- -m 4G -smp 4`.

Options:
  --daemon         Run in the background while keeping the graphical window
  --download-only  Download the base image if it is missing, then exit
  --state-dir DIR  Directory for the VM disk, seed ISO and runtime files
                   (default: .vm/ubuntu2604; overrides UBUNTU_VM_STATE_DIR)
  -h, --help       Show this help

Environment variables:
  UBUNTU_VM_BASE_IMAGE     Base qcow2 image path (default: resolved by shared image helper)
  UBUNTU_VM_SEED_IMAGE     cloud-init seed ISO path (default: alongside the VM disk)
  UBUNTU_VM_STATE_DIR      Runtime state directory (default: .vm/ubuntu2604)
  UBUNTU_VM_DISK_IMAGE     Writable overlay path
  UBUNTU_VM_DISK_SIZE      New overlay virtual capacity (default: 20G)
  VM_IMAGE_ARCH            Guest/host architecture: arm64 or amd64 (default: auto-detected from the host)
  UBUNTU_VM_UEFI_CODE      UEFI read-only firmware path (auto-detected for the selected architecture)
  UBUNTU_VM_UEFI_VARS      UEFI variable-store template path (auto-detected for the selected architecture)
  UBUNTU_VM_DISPLAY        QEMU display backend (default: cocoa on macOS, gtk elsewhere)
  UBUNTU_VM_QEMU_CONFIG    QEMU config template rendered into the state directory and
                           passed to -readconfig (default: qemu/ubuntu2604.ini)
""")


args = sys.argv[1:]
while args:
    arg = args.pop(0)
    if arg == "--daemon":
        daemonize = True
    elif arg == "--download-only":
        download_only = True
    elif arg == "--state-dir":
        if not args:
            die(f"{arg} requires a directory argument.", 2)
        state_dir = Path(args.pop(0))
    elif arg.startswith("--state-dir="):
        state_dir = Path(arg.split("=", 1)[1])
    elif arg in ("-h", "--help"):
        usage()
        sys.exit(0)
    elif arg == "--":
        extra_args.extend(args)
        break
    else:
        print(f"Error: unknown option: {arg}", file=sys.stderr)
        usage(sys.stderr)
        sys.exit(2)

disk_image = Path(env("UBUNTU_VM_DISK_IMAGE") or state_dir / "disk.qcow2")
seed_image = Path(env("UBUNTU_VM_SEED_IMAGE") or disk_image.parent / "seed.iso")
uefi_vars = state_dir / "OVMF_VARS_4M.fd"
pid_file = state_dir / "qemu.pid"
log_file = state_dir / "console.log"
qemu_config = state_dir / "qemu.ini"

if not download_only:
    for command_name in (qemu_binary, "qemu-img"):
        if shutil.which(command_name) is None:
            die(f"{command_name} is not installed.")

if not base_image.is_file():
    if uses_shared_base_image:
        run([str(image_downloader), "ubuntu2604"])
    else:
        die(f"Ubuntu base image was not found: {base_image}")
if not base_image.is_file():
    die(f"downloader did not create: {base_image}")
if download_only:
    print(f"Base image is ready: {base_image}")
    sys.exit(0)
for firmware_file in (uefi_code, uefi_vars_template):
    if not firmware_file.is_file():
        die(f"required UEFI firmware was not found: {firmware_file}")

if not seed_image.is_file():
    seed_builder = script_dir / "seediso" / "ubuntu2604" / "build-seed.sh"
    if not is_executable(seed_builder):
        die(f"seed image is missing and builder is unavailable: {seed_builder}")
    print(f"Seed image not found; building {seed_image}...", flush=True)
    run([str(seed_builder), str(seed_image)])

state_dir.mkdir(parents=True, exist_ok=True)

# QEMU requires the code and vars pflash images to be exactly the same size.
# Homebrew's edk2 firmware ships code/vars pairs at mismatched sizes (true for
# both the aarch64 and x86_64 pairs), so pad per-VM copies of whichever one is
# smaller up to match the other.
uefi_code_flash = state_dir / "OVMF_CODE_4M.fd"
if not uefi_code_flash.is_file() or not uefi_vars.is_file():
    flash_size = max(uefi_code.stat().st_size, uefi_vars_template.stat().st_size)
    for source, target in ((uefi_code, uefi_code_flash), (uefi_vars_template, uefi_vars)):
        if not target.is_file():
            shutil.copyfile(source, target)
            os.truncate(target, flash_size)

if pid_file.is_file():
    old_pid = pid_file.read_text().strip()
    if old_pid.isdigit():
        try:
            os.kill(int(old_pid), 0)
        except OSError:
            pass
        else:
            die(f"the Ubuntu virtual machine is already running (PID {old_pid}).")
    pid_file.unlink(missing_ok=True)

if not disk_image.is_file():
    run(["qemu-img", "create", "-q", "-f", "qcow2", "-F", "qcow2",
         "-b", str(base_image), str(disk_image), disk_size])
    print(f"Created writable disk overlay: {disk_image}")

is_macos = platform.system() == "Darwin"
display_backend = env("UBUNTU_VM_DISPLAY") or ("cocoa" if is_macos else "gtk")

# Most of the VM definition lives in the QEMU config template; fill in the
# file paths and hand the rendered file to QEMU with -readconfig.
config_values = {
    "uefi_code": uefi_code_flash,
    "uefi_vars": uefi_vars,
    "disk_image": disk_image,
    "seed_image": seed_image,
}
for key, value in config_values.items():
    # QEMU's config file parser has no escape syntax for quoted values.
    if any(char in str(value) for char in '"\n'):
        die(f"{key} must not contain double quotes or newlines: {value}")
for config_file in (qemu_config_template, qemu_platform_config):
    if not config_file.is_file():
        die(f"QEMU config file was not found: {config_file}")
try:
    qemu_config_text = string.Template(qemu_config_template.read_text()).substitute(
        {key: str(value) for key, value in config_values.items()})
except (KeyError, ValueError) as error:
    die(f"invalid placeholder in {qemu_config_template}: {error}")
qemu_config.write_text(qemu_config_text)

qemu_args = [
    "-readconfig", str(qemu_config),
    "-readconfig", str(qemu_platform_config),
    # "max" is the host CPU under KVM/HVF and the fullest model under TCG.
    "-cpu", "max",
    # The config file's [netdev] does not stop QEMU from adding its default
    # NIC, a second user network without the SSH port forward.
    "-nic", "none",
    "-display", display_backend,
]

print(f"Starting Ubuntu 26.04 ({host_arch}) from {qemu_config_template} "
      f"and {qemu_platform_config}.", flush=True)
# The SSH port forward lives in the config file, e.g. hostfwd = "tcp::2225-:22".
ssh_forward = re.search(r'hostfwd\s*=\s*"tcp:[^:"]*:(\d+)-[^:"]*:22"', qemu_config_text)
if ssh_forward:
    print(f"SSH is forwarded to localhost:{ssh_forward.group(1)}: "
          f"ssh -p {ssh_forward.group(1)} guofang@localhost", flush=True)

if daemonize:
    log_file.write_text("")
    qemu_args += [
        "-serial", f"file:{log_file}",
        "-monitor", "none",
        "-daemonize",
        "-pidfile", str(pid_file),
    ]
    run([qemu_binary, *qemu_args, *extra_args])
    print(f"VM started (PID {pid_file.read_text().strip()}); console log: {log_file}")
else:
    os.execvp(qemu_binary, [qemu_binary, *qemu_args, *extra_args])