Merge remote-tracking branch 'origin/pr/186'

* origin/pr/186:
  Added support for custom UKI sections, such as ucode

Pull request description:

This enables the addition of sections defined in the [Xen documentation](https://xenbits.xen.org/docs/unstable/misc/efi.html).

For anyone having used the tool before, this differs in the error and help messages, but arguments that worked for the previous version will work with this version as well without modification.

Author: marmarek

uki-generate

@@ -1,17 +1,38 @@
 #!/usr/bin/python3 -I
+import argparse
 import os
 import subprocess
 import sys
 import re
 from typing import List, Tuple
 
+class SectionAdder:
+    def __init__(self, base_addr, round_to_next):
+        self._address = base_addr
+        self._args = []
+        self.round_to_next = round_to_next
+    def add_section(self, section_name, path):
+        self._address = self.round_to_next(self._address)
+        self._args.append(f"--add-section=.{section_name}={path}")
+        self._args.append(f"--change-section-vma=.{section_name}={self._address}")
+        self._address += os.stat(path).st_size
+    def objcopy_args(self):
+        return self._args
+
 def main(args):
     section_re = re.compile(rb"\A *(?:0|[1-9][0-9]*) +([!-~]+) +([0-9a-f]{8}) +([0-9a-f]{16}) +[0-9a-f]{16} +[0-9a-f]{8} +2")
     alignment_mask = (1 << 21) - 1
-    if len(args) != 6:
-        print(f"Usage: uki-generate HYPERVISOR CONFIG KERNEL INITRAMFS OUTPUT", file=sys.stderr)
-        sys.exit(1)
-    _, hyp, cfg, kern, initramfs, out = args
+    parser = argparse.ArgumentParser(
+            prog='uki-generate',
+            description='Create a unified kernel image for Qubes OS')
+    parser.add_argument('hypervisor', help="Xen Kernel")
+    parser.add_argument('config', help="Xen config file")
+    parser.add_argument('kernel', help="Dom0 kernel")
+    parser.add_argument('initramfs')
+    parser.add_argument('output', help="UKI image save path")
+    parser.add_argument('--custom-section', nargs=2, action='append', help='A custom section placed right before the kernel')
+    args = parser.parse_args(args[1:])
+    hyp, cfg, kern, initramfs, custom, out = args.hypervisor, args.config, args.kernel, args.initramfs, args.custom_section, args.output
     if hyp[0] != '/':
         hyp = './' + hyp
     if out[0] != '/':
@@ -38,27 +59,30 @@ def main(args):
             print(f"Fatal error: Address overflow: {f} exceeds {max_address}", file=sys.stderr)
             sys.exit(1)
         return (f + alignment_mask) & ~alignment_mask
-    base_address = round_to_next(max_vma)
-    kernel_vma = round_to_next(base_address + os.stat(cfg).st_size)
-    initramfs_vma = round_to_next(kernel_vma + os.stat(kern).st_size)
-    cmdline = [
-        "objcopy",
-        f"--section-alignment={alignment_mask + 1}",
-        f"--file-alignment={1 << 5}",
-        #"--remove-section=.buildid",
-        "--remove-section=.annobin.*",
-        #"--strip-debug",
-        f"--add-section=.config={cfg}",
-        f"--change-section-vma=.config={base_address}",
-        f"--add-section=.kernel={kern}",
-        f"--change-section-vma=.kernel={kernel_vma}",
-        f"--add-section=.ramdisk={initramfs}",
-        f"--change-section-vma=.ramdisk={initramfs_vma}",
-        "--long-section-names=disable",
-        "--",
-        hyp,
-        out,
-    ]
+
+    adder = SectionAdder(base_addr = max_vma, round_to_next = round_to_next)
+    adder.add_section("config", cfg)
+    for (section, path) in custom or []:
+        adder.add_section(section, path)
+    adder.add_section("kernel", kern)
+    adder.add_section("ramdisk", initramfs)
+    cmdline = (
+        [
+            "objcopy",
+            f"--section-alignment={alignment_mask + 1}",
+            f"--file-alignment={1 << 5}",
+            #"--remove-section=.buildid",
+            "--remove-section=.annobin.*",
+            #"--strip-debug",
+        ]
+        + adder.objcopy_args()
+        + [
+            "--long-section-names=disable",
+            "--",
+            hyp,
+            out,
+        ]
+    )
     subprocess.check_call(cmdline, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
 if __name__ == '__main__':
     try: