org.osbuild.bootiso.mono

#!/usr/bin/python3
import contextlib
import os
import re
import shutil
import subprocess
import sys
import tempfile

import osbuild.api
import osbuild.remoteloop as remoteloop

LORAX_TEMPLATES = "/usr/share/lorax/templates.d"


@contextlib.contextmanager
def mount(source, dest):
    subprocess.run(["mount", source, dest], check=True)
    try:
        yield dest
    finally:
        subprocess.run(["umount", "-R", dest], check=True)


def install(src, dst, mode=None):
    shutil.copyfile(src, dst)
    if mode:
        os.chmod(dst, mode)


def replace(target, patterns):
    finder = [(re.compile(p), s) for p, s in patterns]
    newfile = target + ".replace"

    with open(target, "r", encoding="utf8") as i, open(newfile, "w", encoding="utf8") as o:
        for line in i:
            for p, s in finder:
                line = p.sub(s, line)
            o.write(line)
    os.rename(newfile, target)


def make_rootfs(tree, image, size, workdir, loop_client):
    with open(image, "w", encoding="utf8") as f:
        os.ftruncate(f.fileno(), size)

    root = os.path.join(workdir, "rootfs")
    os.makedirs(root)

    with loop_client.device(image, 0, size) as dev:
        subprocess.run(["mkfs.ext4",
                        "-L", "Anaconda",
                        "-b", "4096",
                        "-m", "0",
                        dev],
                       input="y", encoding='utf8', check=True)

        with mount(dev, root):
            print("copying tree")
            subprocess.run(["cp", "-a", f"{tree}/.", root],
                           check=True)
            print("done")


def make_efi(efi, info, root, loop_client):
    arches = efi["architectures"]
    vendor = efi["vendor"]

    efidir = os.path.join(root, "EFI", "BOOT")
    os.makedirs(efidir)

    # arch related data
    for arch in arches:
        arch = arch.lower()
        targets = [
            (f"shim{arch}.efi", f"BOOT{arch}.EFI".upper()),
            (f"mm{arch}.efi", f"mm{arch}.efi"),
            (f"gcd{arch}.efi", f"grub{arch}.efi")
        ]

        for src, dst in targets:
            shutil.copy2(os.path.join("/boot/efi/EFI/", vendor, src),
                         os.path.join(efidir, dst))

    # the font
    fontdir = os.path.join(efidir, "fonts")
    os.makedirs(fontdir, exist_ok=True)
    shutil.copy2("/usr/share/grub/unicode.pf2", fontdir)

    # the config
    configdir = info["configdir"]
    version = info["version"]
    name = info["name"]
    isolabel = info["isolabel"]
    cmdline = info["cmdline"]

    kdir = "/" + os.path.relpath(info["kerneldir"], start=root)
    print(f"kernel dir at {kdir}")

    config = os.path.join(efidir, "grub.cfg")
    shutil.copy2(os.path.join(configdir, "grub2-efi.cfg"), config)

    replace(config, [
        ("@VERSION@", version),
        ("@PRODUCT@", name),
        ("@KERNELNAME@", "vmlinuz"),
        ("@KERNELPATH@", os.path.join(kdir, "vmlinuz")),
        ("@INITRDPATH@", os.path.join(kdir, "initrd.img")),
        ("@ISOLABEL@", isolabel),
        ("@ROOT@", cmdline)
    ])

    if "IA32" in arches:
        shutil.copy2(config, os.path.join(efidir, "BOOT.cfg"))

    # estimate the size
    blocksize = 2048
    size = blocksize * 256  # blocksize * overhead
    for parent, dirs, files in os.walk(efidir):
        for name in files + dirs:
            t = os.path.join(parent, name)
            s = os.stat(t).st_size
            d = s % blocksize
            if not s or d:
                s += blocksize - d
            size += s
    print(f"Estimates efiboot size to be {size}")

    # create the image
    image = os.path.join(info["imgdir"], "efiboot.img")
    with open(image, "w", encoding="utf8") as f:
        os.ftruncate(f.fileno(), size)

    root = os.path.join(info["workdir"], "mnt")
    os.makedirs(root)

    with loop_client.device(image, 0, size) as dev:
        subprocess.run(["mkfs.fat",
                        "-n", "ANACONDA",
                        dev],
                       input="y", encoding='utf8', check=True)

        with mount(dev, root):
            target = os.path.join(root, "EFI", "BOOT")
            shutil.copytree(efidir, target)
        subprocess.run(["ls", root], check=True)


def make_isolinux(cfg, root, info, tree):
    # the config
    configdir = info["configdir"]
    version = info["version"]
    name = info["name"]
    cmdline = info["cmdline"]
    kerneldir = info["kerneldir"]

    # boot loader
    isolinux = os.path.join(root, "isolinux")
    os.makedirs(isolinux)

    isolinuxfiles = [("isolinux.bin", 0o755),
                     ("ldlinux.c32", 0o755),
                     ("libcom32.c32", 0o755),
                     ("libutil.c32", 0o755),
                     ("vesamenu.c32", 0o755)]
    for target, mode in isolinuxfiles:
        src = os.path.join("/usr/share/syslinux/", target)
        dst = os.path.join(isolinux, target)
        install(src, dst, mode)

    if cfg.get("debug"):
        src = "/usr/share/syslinux/isolinux-debug.bin"
        dst = os.path.join(isolinux, "isolinux.bin")
        install(src, dst, 0o755)

    for target in ["isolinux.cfg", "boot.msg", "grub.conf"]:
        src = os.path.join(configdir, target)
        dst = os.path.join(isolinux, target)
        install(src, dst)

        replace(dst, [
            ("@VERSION@", version),
            ("@PRODUCT@", name),
            ("@ROOT@", cmdline)
        ])

    src = os.path.join(tree, "usr/share/anaconda/boot/syslinux-splash.png")
    dst = os.path.join(isolinux, "splash.png")
    install(src, dst)

    # link the kernel
    os.link(os.path.join(kerneldir, "vmlinuz"),
            os.path.join(isolinux, "vmlinuz"))
    os.link(os.path.join(kerneldir, "initrd.img"),
            os.path.join(isolinux, "initrd.img"))


# pylint: disable=too-many-statements
def main(inputs, root, options, workdir, loop_client):
    tree = inputs["rootfs"]["path"]
    name = options["product"]["name"]
    version = options["product"]["version"]
    kernel = options["kernel"]
    isolabel = options["isolabel"]
    templates = options["templates"]
    efi = options.get("efi")
    isolinux = options.get("isolinux", {})
    kopts = options.get("kernel_opts")
    rootfs = options.get("rootfs", {})

    # input directories
    templatedir = os.path.join(LORAX_TEMPLATES, templates)

    # select the template based on the architecture, where
    # we reuse the efi setting, since we only support efi
    # on aarch64 this is good enough for now
    if efi and "AA64" in efi["architectures"]:
        arch = "aarch64"
    else:
        arch = "x86"

    configdir = os.path.join(templatedir, "config_files", arch)

    # output directories
    imgdir = os.path.join(root, "images")
    pxedir = os.path.join(imgdir, "pxeboot")

    os.makedirs(imgdir)

    # boot configuration
    cmdline = f"inst.stage2=hd:LABEL={isolabel}"
    if kopts:
        cmdline += " " + kopts

    info = {
        "version": version,
        "name": name,
        "isolabel": isolabel,
        "workdir": workdir,
        "configdir": configdir,
        "kerneldir": pxedir,
        "imgdir": imgdir,
        "cmdline": cmdline
    }

    # install the kernel
    kerneldir = pxedir
    kernel_input = inputs.get("kernel", inputs["rootfs"])
    kernel_tree = kernel_input["path"]
    bootdir = os.path.join(kernel_tree, "boot")

    os.makedirs(kerneldir)
    install(os.path.join(bootdir, f"vmlinuz-{kernel}"),
            os.path.join(kerneldir, "vmlinuz"))

    install(os.path.join(bootdir, f"initramfs-{kernel}.img"),
            os.path.join(kerneldir, "initrd.img"))

    # iso linux boot
    if isolinux.get("enabled"):
        make_isolinux(isolinux, root, info, tree)

    # efi boot
    if efi:
        make_efi(efi, info, root, loop_client)

    # install.img
    #   rootfs.img
    liveos_work = os.path.join(workdir, "liveos")
    liveos = os.path.join(liveos_work, "LiveOS")
    os.makedirs(liveos)

    rootfs_size = rootfs.get("size", 3072) * 1024 * 1024
    compression = rootfs.get("compression", {})

    rootfs = os.path.join(liveos, "rootfs.img")
    make_rootfs(tree, rootfs, rootfs_size, workdir, loop_client)

    installimg = os.path.join(imgdir, "install.img")
    cmd = ["mksquashfs", liveos_work, installimg]

    if compression:
        method = compression["method"]
        opts = compression.get("options", {})
        cmd += ["-comp", method]
        for opt, val in opts.items():
            cmd += [f"-X{opt}", val]

    subprocess.run(cmd, check=True)


if __name__ == '__main__':
    args = osbuild.api.arguments()
    _output_dir = args["tree"]
    with tempfile.TemporaryDirectory(dir=_output_dir) as _workdir:
        ret = main(args["inputs"],
                   _output_dir,
                   args["options"],
                   _workdir,
                   remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
    sys.exit(ret)