debian-forge/osbuild

#!/usr/bin/python3

import argparse
import json
import os
import subprocess
import sys
import tempfile


RESET = "\033[0m"
BOLD = "\033[1m"
RED = "\033[31m"


class BuildRoot:
    def __init__(self, path=os.getcwd()):
        self.buildroot = tempfile.mkdtemp(prefix="osbuild-buildroot-", dir=path)
        self.buildroot_mounted = False
        self.tree = tempfile.mkdtemp(prefix="osbuild-tree-", dir=path)
        self.tree_mounted = False
        try:
            subprocess.run(["mount", "-o", "bind,ro", "/", self.buildroot], check=True)
            self.tree_mounted = True
            subprocess.run(["mount", "-t", "tmpfs", "tmpfs", self.tree], check=True)
            self.buildroot_mounted = True
        except subprocess.CalledProcessError:
            self.unmount()
            raise

        # systemd-nspawn silently removes some characters when choosing a
        # machine name from the directory name. The only one relevant for
        # us is '_', because all other characters used by
        # TemporaryDirectory() are allowed. Replace it with 'L's
        # (TemporaryDirectory() only uses lower-case characters)
        self.machine_name = os.path.basename(self.buildroot).replace("_", "L")

    def unmount(self):
        if self.tree:
            if self.tree_mounted:
                subprocess.run(["umount", "--lazy", self.tree], check=True)
            os.rmdir(self.tree)
            self.tree = None
        if self.buildroot:
            if self.buildroot_mounted:
                subprocess.run(["umount", "--lazy", self.buildroot], check=True)
            os.rmdir(self.buildroot)
            self.buildroot = None

    def run(self, argv, binds=[], readonly_binds=[], *args, **kwargs):
        return subprocess.run([
            "systemd-nspawn",
            "--quiet",
            "--as-pid2",
            "--link-journal=no",
            "--volatile=yes",
            f"--machine={self.machine_name}",
            f"--directory={self.buildroot}",
            f"--bind={self.tree}:/tmp/tree",
            *[f"--bind={src}:{dest}" for src, dest in binds],
            *[f"--bind-ro={src}:{dest}" for src, dest in readonly_binds]
        ] + argv, *args, **kwargs)

    def run_stage(self, stage, options={}, input_dir=None, output_dir=None, sit=False):
        options = {
            **options,
            "tree": "/tmp/tree",
            "input_dir": None
        }

        binds = [
            (f"{os.getcwd()}/run-stage", "/tmp/run-stage"),
            (f"{os.getcwd()}/stages/{stage}", "/tmp/stage"),
            ("/etc/pki", "/etc/pki")
        ]

        robinds = []
        if input_dir:
            options["input_dir"] = "/tmp/input"
            robinds.append((input_dir, "/tmp/input"))

        if output_dir:
            options["output_dir"] = "/tmp/output"
            binds.append((output_dir, "/tmp/output"))

        argv = ["/tmp/run-stage"]
        if sit:
            argv.append("--sit")
        argv.append("/tmp/stage")

        self.run(argv, binds=binds, readonly_binds=robinds, input=json.dumps(options), encoding="utf-8", check=True)

    def __del__(self):
        self.unmount()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.unmount()


def run_stage_interactive(i, name, options, buildroot, input_dir=None, output_dir=None, sit=False):
    print()
    print(f"{RESET}{BOLD}{i}. {name}{RESET} " + json.dumps(options or {}, indent=2))
    print("Inspect with:")
    print(f"\t# nsenter -a --wd=/root -t `machinectl show {buildroot.machine_name} -p Leader --value`")
    print()

    try:
        buildroot.run_stage(name, options, input_dir, output_dir, sit)
    except KeyboardInterrupt:
        print()
        print(f"{RESET}{BOLD}{RED}Aborted{RESET}")
        return False
    except subprocess.CalledProcessError as error:
        print()
        print(f"{RESET}{BOLD}{RED}{name} failed with code {error.returncode}{RESET}")
        return False
    else:
        return True


def run_interactive(pipeline_path, input_dir, output_dir, sit):
    if output_dir and len(os.listdir(output_dir)) != 0:
        print()
        print(f"{RESET}{BOLD}{RED}Output directory {output_dir} is not empty{RESET}")
        return False

    with open(pipeline_path) as f:
        pipeline = json.load(f)

    with BuildRoot() as buildroot:
        for i, stage in enumerate(pipeline["stages"], start=1):
            name = stage["name"]
            options = stage.get("options", {})
            if not run_stage_interactive(i, name, options, buildroot, input_dir=input_dir, sit=sit):
                return False

        assembler = pipeline.get("assembler")
        if assembler:
            name = assembler["name"]
            options = assembler.get("options", {})
            if not run_stage_interactive("A", name, options, buildroot, output_dir=output_dir, sit=sit):
                return False

    return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build operating system images")
    parser.add_argument("pipeline_path", metavar="PIPELINE",
                        help="json file containing the pipeline that should be built")
    parser.add_argument("--input", dest="input_dir", metavar="DIRECTORY", type=os.path.abspath,
                        help="provide the contents of DIRECTORY to the first stage")
    parser.add_argument("--output", dest="output_dir", metavar="DIRECTORY", type=os.path.abspath,
                        help="provide the empty DIRECTORY as output argument to the last stage")
    parser.add_argument("--sit", action="store_true",
                        help="keep the build environment up when a stage failed")
    args = parser.parse_args()

    if not run_interactive(args.pipeline_path, args.input_dir, args.output_dir, args.sit):
        sys.exit(1)