nixos/image: add repart builder

1 files changed, 113 insertions, 0 deletions

diff --git a/nixos/modules/image/amend-repart-definitions.py b/nixos/modules/image/amend-repart-definitions.py
new file mode 100644
index 0000000000000..e50ed6fd39a79
--- /dev/null
+++ b/nixos/modules/image/amend-repart-definitions.py
@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+
+"""Amend systemd-repart definiton files.
+
+In order to avoid Import-From-Derivation (IFD) when building images with
+systemd-repart, the definition files created by Nix need to be amended with the
+store paths from the closure.
+
+This is achieved by adding CopyFiles= instructions to the definition files.
+
+The arbitrary files configured via `contents` are also added to the definition
+files using the same mechanism.
+"""
+
+import json
+import sys
+import shutil
+import os
+import tempfile
+from pathlib import Path
+
+
+def add_contents_to_definition(
+    definition: Path, contents: dict[str, dict[str, str]] | None
+) -> None:
+    """Add CopyFiles= instructions to a definition for all files in contents."""
+    if not contents:
+        return
+
+    copy_files_lines: list[str] = []
+    for target, options in contents.items():
+        source = options["source"]
+
+        copy_files_lines.append(f"CopyFiles={source}:{target}\n")
+
+    with open(definition, "a") as f:
+        f.writelines(copy_files_lines)
+
+
+def add_closure_to_definition(
+    definition: Path, closure: Path | None, strip_nix_store_prefix: bool | None
+) -> None:
+    """Add CopyFiles= instructions to a definition for all paths in the closure.
+
+    If strip_nix_store_prefix is True, `/nix/store` is stripped from the target path.
+    """
+    if not closure:
+        return
+
+    copy_files_lines: list[str] = []
+    with open(closure, "r") as f:
+        for line in f:
+            if not isinstance(line, str):
+                continue
+
+            source = Path(line.strip())
+            target = str(source.relative_to("/nix/store/"))
+            target = f":{target}" if strip_nix_store_prefix else ""
+
+            copy_files_lines.append(f"CopyFiles={source}{target}\n")
+
+    with open(definition, "a") as f:
+        f.writelines(copy_files_lines)
+
+
+def main() -> None:
+    """Amend the provided repart definitions by adding CopyFiles= instructions.
+
+    For each file specified in the `contents` field of a partition in the
+    partiton config file, a `CopyFiles=` instruction is added to the
+    corresponding definition file.
+
+    The same is done for every store path of the `closure` field.
+
+    Print the path to a directory that contains the amended repart
+    definitions to stdout.
+    """
+    partition_config_file = sys.argv[1]
+    if not partition_config_file:
+        print("No partition config file was supplied.")
+        sys.exit(1)
+
+    repart_definitions = sys.argv[2]
+    if not repart_definitions:
+        print("No repart definitions were supplied.")
+        sys.exit(1)
+
+    with open(partition_config_file, "rb") as f:
+        partition_config = json.load(f)
+
+    if not partition_config:
+        print("Partition config is empty.")
+        sys.exit(1)
+
+    temp = tempfile.mkdtemp()
+    shutil.copytree(repart_definitions, temp, dirs_exist_ok=True)
+
+    for name, config in partition_config.items():
+        definition = Path(f"{temp}/{name}.conf")
+        os.chmod(definition, 0o644)
+
+        contents = config.get("contents")
+        add_contents_to_definition(definition, contents)
+
+        closure = config.get("closure")
+        strip_nix_store_prefix = config.get("stripStorePaths")
+        add_closure_to_definition(definition, closure, strip_nix_store_prefix)
+
+    print(temp)
+
+
+if __name__ == "__main__":
+    main()