mkboot

#!/usr/bin/python3

import argparse
import contextlib
import os
import pycdlib
import shutil
import subprocess
import sys
import tempfile
from string import Template


@contextlib.contextmanager
def mount_iso(iso, path):

    subprocess.run(["mount", "-t", "iso9660", "-o", "loop,ro", iso, path],
                   check=True)

    try:
        yield path
    finally:
        subprocess.run(["umount", path])


def delete_generated(path, verbose):
    for root, dirs, files in os.walk(path):
        for f in files:
            if f in ["TRANS.TBL", "boot.cat", "boot.catalog"]:
                t = os.path.join(root, f)
                if verbose:
                    print(f"unlink {t}")
                os.unlink(t)


def make_iso(path, label, output, workdir, verbose):
    cmd = [
        "/usr/bin/xorrisofs"
    ]

    if verbose:
        cmd += ["-verbose"]
    else:
        cmd += ["-quiet"]

    cmd += [
        "-V", label,
    ]

    cmd += [
        "-b", "isolinux/isolinux.bin",
        "-c", "isolinux/boot.cat",
        "-no-emul-boot",
        "-boot-load-size", "4",
        "-boot-info-table"
    ]

    cmd += [
        "-rock", "-joliet"
    ]

    cmd += [
        "-eltorito-alt-boot",
        "-e", "images/efiboot.img",
        "-no-emul-boot",
        "-isohybrid-gpt-basdat",
    ]

    if os.path.exists("/usr/share/syslinux/isohdpfx.bin"):
        cmd += [
            "-isohybrid-mbr", "/usr/share/syslinux/isohdpfx.bin",
        ]

    cmd += [
        '-o', output,
        path
    ]

    if verbose:
        print(cmd)
        stdout = sys.stdout
    else:
        stdout = subprocess.DEVNULL

    subprocess.run(cmd,
                   cwd=workdir,
                   stdout=stdout,
                   stderr=stdout,
                   check=True)


def get_label(path):
    try:
        iso = pycdlib.PyCdlib()
        iso.open(path)
        label = iso.pvd.volume_identifier.decode("UTF-8").strip()
    except pycdlib.PyCdlibException as e:
        label = ""
    return label


def write_template(source, dest, **subst):
    with open(source, "r") as f:
        tpl = Template(f.read())
        data = tpl.substitute(**subst)
    with open(dest, "w") as f:
        f.write(data)


def main():
    parser = argparse.ArgumentParser(description="Modify boot iso")
    parser.add_argument("iso", metavar="ISO", type=os.path.abspath,
                        help="original ISO to modify")
    parser.add_argument("--output", metavar="FILENAME", type=os.path.abspath,
                        default="bootiso.iso",
                        help="Name of the output file")
    parser.add_argument("--kickstart", metavar="FILENAME", type=os.path.abspath,
                        default="edge.ks",
                        help="Name of the kickstart file to embed")
    parser.add_argument("--verbose", default=False, action="store_true",
                        help="Show more information")
    parser.add_argument("--kargs", metavar="ID", action="append", type=str, default=["quiet"],
                        help="Additional kernel command line args")
    args = parser.parse_args(sys.argv[1:])

    bootiso = args.iso
    output = args.output
    ks_path = args.kickstart
    ks_name = os.path.basename(ks_path)
    kargs = " ".join(args.kargs)

    label = get_label(bootiso)
    print(f"Label:     {label}")
    print(f"Kickstart: {ks_name} ({ks_path})")
    print(f"Kargs:     {kargs}")

    with tempfile.TemporaryDirectory(dir="/var/tmp") as tmp:
        mountpath = os.path.join(tmp, "mount")
        isopath = os.path.join(tmp, "iso")
        os.makedirs(mountpath)
        os.makedirs(isopath)

        with mount_iso(bootiso, mountpath) as path:
            subprocess.run(["cp", "-a", f"{path}/.", f"{isopath}/"],
                           check=True)

            delete_generated(isopath, args.verbose)

            bootcfg = {
                "isolinux.cfg" : "isolinux",
                "grub.cfg": "EFI/BOOT",
            }

            subst = {
                "label": label,
                "kickstart": ks_name,
                "kargs": kargs
            }

            for filename, path in bootcfg.items():
                src = os.path.join("boot", filename)
                dst = os.path.join(isopath, path, filename)
                write_template(src, dst, **subst)

            # write the kickstart
            dst = os.path.join(isopath, ks_name)
            write_template(ks_path, dst, **subst)

            make_iso(isopath, label, output, tmp, args.verbose)


if __name__ == "__main__":
    main()