Refactor copy_{uki,vmlinuz,initrd}

A follow-up commit will introduce the ability to disable copying these
to the output directory, so refactor all the logic so that they are
contained within their respectiv functions.

mkosi/__init__.py

@@ -2152,12 +2152,9 @@ def copy_nspawn_settings(context: Context) -> None:
         shutil.copy2(context.config.nspawn_settings, context.staging / context.config.output_nspawn_settings)
 
 
-def copy_uki(context: Context) -> None:
-    if (context.staging / context.config.output_split_uki).exists():
-        return
-
+def get_uki(context: Context) -> Optional[Path]:
     if not want_efi(context.config) or context.config.unified_kernel_images == ConfigFeature.disabled:
-        return
+        return None
 
     ukis = sorted(
         (context.root / "boot/EFI/Linux").glob("*.efi"),
@@ -2168,30 +2165,36 @@ def copy_uki(context: Context) -> None:
     if (uki := context.root / efi_boot_binary(context)).exists() and (
         KernelType.identify(context.config, uki) == KernelType.uki
     ):
-        pass
+        return uki
     elif (uki := context.root / shim_second_stage_binary(context)).exists() and (
         KernelType.identify(context.config, uki) == KernelType.uki
     ):
-        pass
-    elif ukis:
-        uki = ukis[0]
-    else:
-        return
+        return uki
 
-    shutil.copy(uki, context.staging / context.config.output_split_uki)
+    if not ukis:
+        return None
 
-    # Extract the combined initrds from the UKI so we can use it to direct kernel boot with qemu if needed.
-    extract_pe_section(context, uki, ".initrd", context.staging / context.config.output_split_initrd)
+    return ukis[0]
 
-    # ukify will have signed the kernel image as well. Let's make sure we put the signed kernel
-    # image in the output directory instead of the unsigned one by reading it from the UKI.
-    extract_pe_section(context, uki, ".linux", context.staging / context.config.output_split_kernel)
+
+def copy_uki(context: Context) -> None:
+    if (context.staging / context.config.output_split_uki).exists():
+        return
+
+    if uki := get_uki(context):
+        shutil.copy(uki, context.staging / context.config.output_split_uki)
 
 
 def copy_vmlinuz(context: Context) -> None:
     if (context.staging / context.config.output_split_kernel).exists():
         return
 
+    # ukify will have signed the kernel image as well. Let's make sure we put the signed kernel
+    # image in the output directory instead of the unsigned one by reading it from the UKI.
+    if uki := get_uki(context):
+        extract_pe_section(context, uki, ".linux", context.staging / context.config.output_split_kernel)
+        return
+
     for _, kimg in gen_kernel_images(context):
         shutil.copy(context.root / kimg, context.staging / context.config.output_split_kernel)
         break
@@ -2204,6 +2207,11 @@ def copy_initrd(context: Context) -> None:
     if (context.staging / context.config.output_split_initrd).exists():
         return
 
+    # Extract the combined initrds from the UKI so we can use it to direct kernel boot with qemu if needed.
+    if uki := get_uki(context):
+        extract_pe_section(context, uki, ".initrd", context.staging / context.config.output_split_initrd)
+        return
+
     for kver, _ in gen_kernel_images(context):
         initrds = finalize_initrds(context)