nixos/modules/image/amend-repart-definitions.py

   1 #!/usr/bin/env python
   2
   3 """Amend systemd-repart definiton files.
   4
   5 In order to avoid Import-From-Derivation (IFD) when building images with
   6 systemd-repart, the definition files created by Nix need to be amended with the
   7 store paths from the closure.
   8
   9 This is achieved by adding CopyFiles= instructions to the definition files.
  10
  11 The arbitrary files configured via `contents` are also added to the definition
  12 files using the same mechanism.
  13 """
  14
  15 import json
  16 import sys
  17 import shutil
  18 from pathlib import Path
  19
  20
  21 def add_contents_to_definition(
  22     definition: Path, contents: dict[str, dict[str, str]] | None
  23 ) -> None:
  24     """Add CopyFiles= instructions to a definition for all files in contents."""
  25     if not contents:
  26         return
  27
  28     copy_files_lines: list[str] = []
  29     for target, options in contents.items():
  30         source = options["source"]
  31
  32         copy_files_lines.append(f"CopyFiles={source}:{target}\n")
  33
  34     with open(definition, "a") as f:
  35         f.writelines(copy_files_lines)
  36
  37
  38 def add_closure_to_definition(
  39     definition: Path, closure: Path | None, strip_nix_store_prefix: bool | None
  40 ) -> None:
  41     """Add CopyFiles= instructions to a definition for all paths in the closure.
  42
  43     If strip_nix_store_prefix is True, `/nix/store` is stripped from the target path.
  44     """
  45     if not closure:
  46         return
  47
  48     copy_files_lines: list[str] = []
  49     with open(closure, "r") as f:
  50         for line in f:
  51             if not isinstance(line, str):
  52                 continue
  53
  54             source = Path(line.strip())
  55             target = str(source.relative_to("/nix/store/"))
  56             target = f":/{target}" if strip_nix_store_prefix else ""
  57
  58             copy_files_lines.append(f"CopyFiles={source}{target}\n")
  59
  60     with open(definition, "a") as f:
  61         f.writelines(copy_files_lines)
  62
  63
  64 def main() -> None:
  65     """Amend the provided repart definitions by adding CopyFiles= instructions.
  66
  67     For each file specified in the `contents` field of a partition in the
  68     partiton config file, a `CopyFiles=` instruction is added to the
  69     corresponding definition file.
  70
  71     The same is done for every store path of the `closure` field.
  72
  73     Print the path to a directory that contains the amended repart
  74     definitions to stdout.
  75     """
  76     partition_config_file = sys.argv[1]
  77     if not partition_config_file:
  78         print("No partition config file was supplied.")
  79         sys.exit(1)
  80
  81     repart_definitions = sys.argv[2]
  82     if not repart_definitions:
  83         print("No repart definitions were supplied.")
  84         sys.exit(1)
  85
  86     with open(partition_config_file, "rb") as f:
  87         partition_config = json.load(f)
  88
  89     if not partition_config:
  90         print("Partition config is empty.")
  91         sys.exit(1)
  92
  93     target_dir = Path("amended-repart.d")
  94     target_dir.mkdir()
  95     shutil.copytree(repart_definitions, target_dir, dirs_exist_ok=True)
  96
  97     for name, config in partition_config.items():
  98         definition = target_dir.joinpath(f"{name}.conf")
  99         definition.chmod(0o644)
 100
 101         contents = config.get("contents")
 102         add_contents_to_definition(definition, contents)
 103
 104         closure = config.get("closure")
 105         strip_nix_store_prefix = config.get("stripNixStorePrefix")
 106         add_closure_to_definition(definition, closure, strip_nix_store_prefix)
 107
 108     print(target_dir.absolute())
 109
 110
 111 if __name__ == "__main__":
 112     main()