debian-forge/assemblers/org.osbuild.qemu

#!/usr/bin/python3

import contextlib
import json
import os
import socket
import shutil
import subprocess
import sys
import tempfile
import osbuild.remoteloop as remoteloop

STAGE_DESC = "Assemble a bootable partitioned disk image with qemu-img"
STAGE_INFO = """
Assemble a bootable partitioned disk image using `qemu-img`.

Creates a sparse MBR-partitioned disk image of the given `size`, with a single
bootable partition containing the root filesystem.

Installs GRUB2 (using the buildhost's `/usr/lib/grub/i386-pc/boot.img` etc.) as
the bootloader.

Copies the tree contents into the root filesystem and then converts the raw
sparse image into the format requested with the `fmt` option.

Buildhost commands used: `truncate`, `mount`, `umount`, `sfdisk`,
`grub2-mkimage`, `mkfs.ext4` or `mkfs.xfs`, `qemu-img`.
"""
STAGE_OPTS = """
"required": ["format", "filename", "ptuuid", "root_fs_uuid", "size"],
"properties": {
  "format": {
    "description": "Image file format to use",
    "type": "string",
    "enum": ["raw", "qcow2", "vdi", "vmdk", "vpc"]
  },
  "filename": {
    "description": "Image filename",
    "type": "string"
  },
  "ptuuid": {
    "description": "UUID for the disk image's partition table",
    "type": "string"
  },
  "root_fs_uuid": {
    "description": "UUID for the root filesystem",
    "type": "string"
  },
  "size": {
    "description": "Virtual disk size",
    "type": "string"
  },
  "root_fs_type": {
    "description": "Type of the root filesystem",
    "type": "string",
    "enum": ["ext4", "xfs"],
    "default": "ext4"
  }
}
"""

@contextlib.contextmanager
def mount(source):
    with tempfile.TemporaryDirectory(prefix="osbuild-mnt") as dest:
        subprocess.run(["mount", source, dest], check=True)
        try:
            yield dest
        finally:
            subprocess.run(["umount", "-R", dest], check=True)


def mkfs_ext4(device, uuid):
    subprocess.run(["mkfs.ext4", "-U", uuid, device], input="y", encoding='utf-8', check=True)


def mkfs_xfs(device, uuid):
    subprocess.run(["mkfs.xfs", "-m", f"uuid={uuid}", device], encoding='utf-8', check=True)


def main(tree, output_dir, options, loop_client):
    fmt = options["format"]
    filename = options["filename"]
    ptuuid = options["ptuuid"]
    root_fs_uuid = options["root_fs_uuid"]
    size = options["size"]
    root_fs_type = options.get("root_fs_type", "ext4")

    # sfdisk works on sectors of 512 bytes and ignores excess space - be explicit about this
    if size % 512 != 0:
        raise ValueError("`size` must be a multiple of sector size (512)")

    if fmt not in ["raw", "qcow2", "vdi", "vmdk", "vpc"]:
        raise ValueError("`format` must be one of raw, qcow, vdi, vmdk, vpc")

    if root_fs_type == "ext4":
        mkfs = mkfs_ext4
        grub2_fs_module = "ext2"
    elif root_fs_type == "xfs":
        mkfs = mkfs_xfs
        grub2_fs_module = "xfs"
    else:
        raise ValueError("`root_fs_type` must be either ext4 or xfs")

    image = "/var/tmp/osbuild-image.raw"
    grub2_core = "/var/tmp/grub2-core.img"

    # Create an empty image file
    subprocess.run(["truncate", "--size", str(size), image], check=True)

    # Set up the partition table of the image
    partition_table = f"label: mbr\nlabel-id: {ptuuid}\nbootable, type=83"
    subprocess.run(["sfdisk", "-q", image], input=partition_table, encoding='utf-8', check=True)

    r = subprocess.run(["sfdisk", "--json", image], stdout=subprocess.PIPE, encoding='utf-8', check=True)
    partition_table = json.loads(r.stdout)
    partition = partition_table["partitiontable"]["partitions"][0]
    partition_offset = partition["start"] * 512
    partition_size = partition["size"] * 512

    # Create the level-2 bootloader
    # The purpose of this is to find the grub modules and configuration
    # to be able to start the level-3 bootloader. It contains the modules
    # necessary to do this, but nothing else.
    subprocess.run(["grub2-mkimage",
                    "--verbose",
                    "--directory", "/usr/lib/grub/i386-pc",
                    "--prefix", "(,msdos1)/boot/grub2",
                    "--format", "i386-pc",
                    "--compression", "auto",
                    "--output", grub2_core,
                    "part_msdos", grub2_fs_module, "biosdisk"],
                   check=True)

    assert os.path.getsize(grub2_core) < partition_offset - 512

    with open(image, "rb+") as image_f:
        # Install the level-1 bootloader into the start of the MBR
        # The purpose of this is simply to jump into the level-2 bootloader.
        with open("/usr/lib/grub/i386-pc/boot.img", "rb") as boot_f:
            # The boot.img file is 512 bytes, but we must only copy the first 440
            # bytes, as these contain the bootstrapping code. The rest of the
            # first sector contains the partition table, and must not be
            # overwritten.
            image_f.write(boot_f.read(440))

        # Install the level-2 bootloader into the space after the MBR, before
        # the first partition.
        with open(grub2_core, "rb") as core_f:
            image_f.seek(512)
            shutil.copyfileobj(core_f, image_f)

    with loop_client.device(image, partition_offset, partition_size) as loop:
        # Populate the first partition of the image with a filesystem
        mkfs(loop, root_fs_uuid)
        # Copy the tree into the target image
        with mount(loop) as mountpoint:
            subprocess.run(["cp", "-a", f"{tree}/.", mountpoint], check=True)

    extra_args = []

    # raw and vdi don't suppport compression
    if fmt not in ("raw", "vdi", "vpc"):
        extra_args.append("-c")

    if fmt == "vpc":
        extra_args += ["-o", "subformat=fixed,force_size"]

    subprocess.run(["qemu-img", "convert", "-O", fmt, *extra_args, image, f"{output_dir}/{filename}"], check=True)


if __name__ == '__main__':
    args = json.load(sys.stdin)

    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
        sock.connect("/run/osbuild/api/remoteloop")
        ret = main(args["tree"], args["output_dir"], args["options"], remoteloop.LoopClient(sock))

    sys.exit(ret)