ukify: add new --extend= switch for importing an existing UKI's sections to later extend

This options is pretty simple, it allows specifying an UKI whose
sections to import first, and place at the beginning of the new UKI.

This is useful for generating multi-profile UKIs piecemeal: generate the
base UKI first, then append a profile, and another one and another one.

The sections imported this way are not included in any PCR signature,
the assumption is that that already happened before in the imported UKI.

2 files changed, 119 insertions, 17 deletions

diff --git a/man/ukify.xml b/man/ukify.xml
index 1c52c9414b..5e7681279b 100644
--- a/man/ukify.xml
+++ b/man/ukify.xml
@@ -229,6 +229,18 @@
         </varlistentry>
 
         <varlistentry>
+          <term><option>--extend=<replaceable>PATH</replaceable></option></term>
+
+          <listitem><para>Takes a path to an existing PE file to import into the newly generated PE file. If
+          this option is used all UKI PE sections of the specified PE file are copied into the target PE file
+          before any new PE sections are appended. This is useful for generating multi-profile UKIs. Note
+          that this only copies PE sections that are defined by the UKI specification, and ignores any other,
+          for example <literal>.text</literal> or similar.</para>
+
+          <xi:include href="version-info.xml" xpointer="v257"/></listitem>
+        </varlistentry>
+
+        <varlistentry>
           <term><option>--tools=<replaceable>DIRS</replaceable></option></term>
 
           <listitem><para>Specify one or more directories with helper tools. <command>ukify</command> will
@@ -703,6 +715,48 @@ Writing public key for PCR signing to /etc/kernel/pcr-system.pub.pem
       by default, so after this file has been created, installations of kernels that create a UKI on the
       local machine using <command>kernel-install</command> will perform signing using this config.</para>
     </example>
+
+    <example>
+      <title>Multi-Profile PE</title>
+
+      <para>First, create the base UKI:</para>
+      <programlisting>$ ukify build \
+      --linux=/lib/modules/6.0.9-300.fc37.x86_64/vmlinuz \
+      --initrd=/some/path/initramfs-6.0.9-300.fc37.x86_64.img \
+      --cmdline='quiet rw' \
+      --output=base.efi
+      </programlisting>
+
+      <para>Then, extend the base UKI with information about profile @0:</para>
+
+      <programlisting>$ ukify build \
+      --extend=base.efi \
+      --profile='TITLE=Base' \
+      --output=base-with-profile-0.efi
+      </programlisting>
+
+      <para>Add a second profile (@1):</para>
+
+      <programlisting>$ ukify build \
+      --extend=base-with-profile-0.efi \
+      --profile='TITLE=Boot into Storage Target Mode
+ID=storagetm' \
+      --cmdline='quiet rw rd.systemd.unit=stroage-target-mode.target' \
+      --output=base-with-profile-0-1.efi
+      </programlisting>
+
+      <para>Add a third profile (@2):</para>
+
+      <programlisting>$ ukify build \
+      --extend=base-with-profile-0-1.efi \
+      --profile='TITLE=Factory Reset
+ID=factory-reset' \
+      --cmdline='quiet rw systemd.unit=factory-reset.target' \
+      --output=base-with-profile-0-1-2.efi
+      </programlisting>
+
+      <para>The resulting UKI <filename>base-with-profile-0-1-2.efi</filename> will now contain three profiles.</para>
+    </example>
   </refsect1>
 
   <refsect1>
diff --git a/src/ukify/ukify.py b/src/ukify/ukify.py
index c8b6ce401b..b3fab755b6 100755
--- a/src/ukify/ukify.py
+++ b/src/ukify/ukify.py
@@ -381,7 +381,14 @@ class UKI:
     sections: list[Section] = dataclasses.field(default_factory=list, init=False)
 
     def add_section(self, section):
-        if section.name in [s.name for s in self.sections]:
+        start = 0
+
+        # Start search at last .profile section, if there is one
+        for i in range(len(self.sections)):
+            if self.sections[i].name == ".profile":
+                start = i+1
+
+        if section.name in [s.name for s in self.sections[start:]]:
             raise ValueError(f'Duplicate section {section.name}')
 
         self.sections += [section]
@@ -493,6 +500,10 @@ def key_path_groups(opts):
                    pp_groups)
 
 
+def pe_strip_section_name(name):
+    return name.rstrip(b"\x00").decode()
+
+
 def call_systemd_measure(uki, linux, opts):
     measure_tool = find_tool('systemd-measure',
                              '/usr/lib/systemd/systemd-measure',
@@ -632,6 +643,9 @@ def pe_add_sections(uki: UKI, output: str):
         # We could strip the signatures, but why would anyone sign the stub?
         raise PEError('Stub image is signed, refusing.')
 
+    # Remember how many sections originate from systemd-stub
+    n_original_sections = len(pe.sections)
+
     for section in uki.sections:
         new_section = pefile.SectionStructure(pe.__IMAGE_SECTION_HEADER_format__, pe=pe)
         new_section.__unpack__(b'\0' * new_section.sizeof())
@@ -665,8 +679,8 @@ def pe_add_sections(uki: UKI, output: str):
         # Special case, mostly for .sbat: the stub will already have a .sbat section, but we want to append
         # the one from the kernel to it. It should be small enough to fit in the existing section, so just
         # swap the data.
-        for i, s in enumerate(pe.sections):
-            if s.Name.rstrip(b"\x00").decode() == section.name:
+        for i, s in enumerate(pe.sections[:n_original_sections]):
+            if pe_strip_section_name(s.Name) == section.name:
                 if new_section.Misc_VirtualSize > s.SizeOfRawData:
                     raise PEError(f'Not enough space in existing section {section.name} to append new data.')
 
@@ -702,7 +716,7 @@ def merge_sbat(input_pe: [pathlib.Path], input_text: [str]) -> str:
             continue
 
         for section in pe.sections:
-            if section.Name.rstrip(b"\x00").decode() == ".sbat":
+            if pe_strip_section_name(section.Name) == ".sbat":
                 split = section.get_data().rstrip(b"\x00").decode().splitlines()
                 if not split[0].startswith('sbat,'):
                     print(f"{f} does not contain a valid SBAT section, skipping.")
@@ -784,6 +798,28 @@ def verify(tool, opts):
 
     return tool['output'] in info
 
+
+def import_to_extend(uki, opts):
+
+    if opts.extend is None:
+        return
+
+    import_sections = ('.linux', '.osrel', '.cmdline', '.initrd',
+                       '.ucode', '.splash', '.dtb', '.uname',
+                       '.sbat', '.pcrsig', '.pcrpkey', '.profile')
+
+    pe = pefile.PE(opts.extend, fast_load=True)
+
+    for section in pe.sections:
+        n = pe_strip_section_name(section.Name)
+
+        if n not in import_sections:
+            continue
+
+        print(f"Copying section '{n}' from '{opts.extend}': {section.Misc_VirtualSize} bytes")
+        uki.add_section(Section.create(n, section.get_data(length=section.Misc_VirtualSize), measure=False))
+
+
 def make_uki(opts):
     # kernel payload signing
 
@@ -848,6 +884,9 @@ def make_uki(opts):
                 format=serialization.PublicFormat.SubjectPublicKeyInfo,
             )
 
+    # Import an existing UKI for extension
+    import_to_extend(uki, opts)
+
     sections = [
         # name,      content,         measure?
         ('.profile', opts.profile,    True ),
@@ -872,21 +911,22 @@ def make_uki(opts):
     for section in opts.sections:
         uki.add_section(section)
 
-    if linux is not None:
-        # Merge the .sbat sections from stub, kernel and parameter, so that revocation can be done on either.
-        input_pes = [opts.stub, linux]
-        if not opts.sbat:
-            opts.sbat = ["""sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
+    if opts.extend is None:
+        if linux is not None:
+            # Merge the .sbat sections from stub, kernel and parameter, so that revocation can be done on either.
+            input_pes = [opts.stub, linux]
+            if not opts.sbat:
+                opts.sbat = ["""sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
 uki,1,UKI,uki,1,https://uapi-group.org/specifications/specs/unified_kernel_image/
 """]
-    else:
-        # Addons don't use the stub so we add SBAT manually
-        input_pes = []
-        if not opts.sbat:
-            opts.sbat = ["""sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
+        else:
+            # Addons don't use the stub so we add SBAT manually
+            input_pes = []
+            if not opts.sbat:
+                opts.sbat = ["""sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
 uki-addon,1,UKI Addon,addon,1,https://www.freedesktop.org/software/systemd/man/latest/systemd-stub.html
 """]
-    uki.add_section(Section.create('.sbat', merge_sbat(input_pes, opts.sbat), measure=linux is not None))
+        uki.add_section(Section.create('.sbat', merge_sbat(input_pes, opts.sbat), measure=linux is not None))
 
     # PCR measurement and signing
 
@@ -1043,7 +1083,7 @@ def generate_keys(opts):
 
 
 def inspect_section(opts, section):
-    name = section.Name.rstrip(b"\x00").decode()
+    name = pe_strip_section_name(section.Name)
 
     # find the config for this section in opts and whether to show it
     config = opts.sections_by_name.get(name, None)
@@ -1384,6 +1424,14 @@ CONFIG_ITEMS = [
     ),
 
     ConfigItem(
+        '--extend',
+        metavar = 'UKI',
+        type = pathlib.Path,
+        help = 'path to existing UKI file whose relevant sections to insert into the UKI first',
+        config_key = 'UKI/Extend',
+    ),
+
+    ConfigItem(
         '--pcr-banks',
         metavar = 'BANK…',
         type = parse_banks,
@@ -1687,7 +1735,7 @@ def finalize_options(opts):
         opts.efi_arch = guess_efi_arch()
 
     if opts.stub is None:
-        if opts.linux is not None:
+        if opts.linux is not None or opts.extend is not None:
             opts.stub = pathlib.Path(f'/usr/lib/systemd/boot/efi/linux{opts.efi_arch}.efi.stub')
         else:
             opts.stub = pathlib.Path(f'/usr/lib/systemd/boot/efi/addon{opts.efi_arch}.efi.stub')