osbuild.py: add BuildRoot.run_assembler()

This is the canonical way to run an assembler.

Also improve error handling by introducing a StageFailed exception.

osbuild

@@ -13,33 +13,7 @@ BOLD = "\033[1m"
 RED = "\033[31m"
 
 
-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)
 
@@ -47,17 +21,23 @@ def run_interactive(pipeline_path, input_dir, output_dir, sit):
         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
+            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()
+            buildroot.run_stage(name, options, input_dir, sit)
 
         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
+            print()
+            print(f"{RESET}{BOLD}Assembling: {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()
+            buildroot.run_assembler(name, options, output_dir)
 
 
 if __name__ == "__main__":
@@ -74,5 +54,17 @@ if __name__ == "__main__":
 
     os.makedirs("/run/osbuild", exist_ok=True)
 
-    if not run_interactive(args.pipeline_path, args.input_dir, args.output_dir, args.sit):
+    if args.output_dir and os.path.exists(args.output_dir) and len(os.listdir(args.output_dir)) != 0:
+        print("Error: output directory is not empty", file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        run_interactive(args.pipeline_path, args.input_dir, args.output_dir, args.sit)
+    except KeyboardInterrupt:
+        print()
+        print(f"{RESET}{BOLD}{RED}Aborted{RESET}")
+        sys.exit(130)
+    except osbuild.StageFailed as error:
+        print()
+        print(f"{RESET}{BOLD}{RED}{error.stage} failed with code {error.returncode}{RESET}")
         sys.exit(1)

osbuild.py

@@ -6,7 +6,10 @@ import sys
 import tempfile
 
 
-__all__ = [ "BuildRoot" ]
+__all__ = [
+    "StageFailed",
+    "BuildRoot"
+]
 
 
 libdir = os.path.dirname(__file__)
@@ -14,6 +17,12 @@ if not os.path.exists(f"{libdir}/stages"):
     libdir = f"{sys.prefix}/lib"
 
 
+class StageFailed(Exception):
+    def __init__(self, stage, returncode):
+        self.stage = stage
+        self.returncode = returncode
+
+
 class BuildRoot:
     def __init__(self, path=None):
         self.buildroot = tempfile.mkdtemp(prefix="osbuild-buildroot-", dir=path)
@@ -62,7 +71,7 @@ class BuildRoot:
             *[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):
+    def run_stage(self, stage, options={}, input_dir=None, sit=False):
         options = {
             **options,
             "tree": "/tmp/tree",
@@ -80,6 +89,32 @@ class BuildRoot:
             options["input_dir"] = "/tmp/input"
             robinds.append((input_dir, "/tmp/input"))
 
+        argv = ["/tmp/run-stage"]
+        if sit:
+            argv.append("--sit")
+        argv.append("/tmp/stage")
+
+        try:
+            self.run(argv, binds=binds, readonly_binds=robinds, input=json.dumps(options), encoding="utf-8", check=True)
+        except subprocess.CalledProcessError as error:
+            raise StageFailed(stage, error.returncode)
+
+    def run_assembler(self, name, options, output_dir=None, sit=False):
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        options = {
+            **options,
+            "tree": "/tmp/tree",
+            "input_dir": None
+        }
+
+        binds = [
+            (f"{libdir}/run-stage", "/tmp/run-stage"),
+            (f"{libdir}/stages/{name}", "/tmp/stage"),
+            ("/etc/pki", "/etc/pki")
+        ]
+
         if output_dir:
             options["output_dir"] = "/tmp/output"
             binds.append((output_dir, "/tmp/output"))
@@ -89,7 +124,10 @@ class BuildRoot:
             argv.append("--sit")
         argv.append("/tmp/stage")
 
-        self.run(argv, binds=binds, readonly_binds=robinds, input=json.dumps(options), encoding="utf-8", check=True)
+        try:
+            self.run(argv, binds=binds, input=json.dumps(options), encoding="utf-8", check=True)
+        except subprocess.CalledProcessError as error:
+            raise StageFailed(stage, error.returncode)
 
     def __del__(self):
         self.unmount()