From 41b5f830c6d74dca69cfea1c259f7e21d4ad9bb9 Mon Sep 17 00:00:00 2001
From: Alejandro Sirgo Rica <asirgo@soleta.eu>
Date: Mon, 22 Apr 2024 13:55:21 +0200
Subject: utils: consolidate code to find efi loader

Add find_windows_efi_loader and find_linux_efi_loader to reduce
code duplication and to centralize efi loader path modifications.
---
 src/utils/boot.py | 26 ++++++--------------------
 src/utils/uefi.py | 33 +++++++++++++++++++++++----------
 2 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/src/utils/boot.py b/src/utils/boot.py
index 399f4ba..244f994 100644
--- a/src/utils/boot.py
+++ b/src/utils/boot.py
@@ -62,19 +62,12 @@ def _boot_uefi_windows(disk, part, mountpoint):
     if not mount_mkdir(esp, esp_mountpoint):
         raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')
 
-    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/bootmgfw.efi',
-                    f'{esp_mountpoint}/EFI/Microsoft/Boot/bootmgfw.efi']
     try:
-        for efi_app in loader_paths:
-            if os.path.exists(efi_app):
-                loader = efi_app[len(esp_mountpoint):]
-                logging.info(f'Found bootloader at ESP partition: {loader}')
-                break
-        else:
-            raise OgError(f'Unable to locate Windows EFI bootloader bootmgfw.efi')
+        loader_path = find_windows_efi_loader(esp_mountpoint, bootlabel)
+        loader_path = loader_path[len(esp_mountpoint):]
 
         efibootmgr_delete_bootentry(bootlabel)
-        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader, bootlabel)
+        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader_path, bootlabel)
         efibootmgr_bootnext(bootlabel)
     finally:
         umount(esp_mountpoint)
@@ -87,19 +80,12 @@ def _boot_uefi_linux(disk, part, mountpoint):
     if not mount_mkdir(esp, esp_mountpoint):
         raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')
 
-    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/shimx64.efi',
-                    f'{esp_mountpoint}/EFI/ubuntu/shimx64.efi']
     try:
-        for efi_app in loader_paths:
-            if os.path.exists(efi_app):
-                loader = efi_app[len(esp_mountpoint):]
-                logging.info(f'Found bootloader at ESP partition: {loader}')
-                break
-        else:
-            raise OgError(f'Unable to locate Linux EFI bootloader shimx64.efi')
+        loader_path = find_linux_efi_loader(esp_mountpoint, bootlabel)
+        loader_path = loader_path[len(esp_mountpoint):]
 
         efibootmgr_delete_bootentry(bootlabel)
-        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader, bootlabel)
+        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader_path, bootlabel)
         efibootmgr_bootnext(bootlabel)
     finally:
         umount(esp_mountpoint)
diff --git a/src/utils/uefi.py b/src/utils/uefi.py
index 0906ae9..48fb6c3 100644
--- a/src/utils/uefi.py
+++ b/src/utils/uefi.py
@@ -103,6 +103,28 @@ def efibootmgr_create_bootentry(disk, part, loader, label, add_to_bootorder=True
         raise OgError(f'Unexpected error adding boot entry to nvram. UEFI firmware might be buggy') from e
 
 
+def _find_efi_loader(loader_paths):
+    for efi_app in loader_paths:
+        logging.info(f'Searching for {efi_app}')
+        if os.path.exists(efi_app):
+            logging.info(f'Found bootloader at ESP partition: {efi_app}')
+            return efi_app
+    else:
+        raise OgError(f'Unable to locate Windows EFI bootloader bootmgfw.efi')
+
+
+def find_windows_efi_loader(esp_mountpoint, bootlabel):
+    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/bootmgfw.efi',
+                    f'{esp_mountpoint}/EFI/Microsoft/Boot/bootmgfw.efi']
+    return _find_efi_loader(loader_paths)
+
+
+def find_linux_efi_loader(esp_mountpoint, bootlabel):
+    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/shimx64.efi',
+                    f'{esp_mountpoint}/EFI/ubuntu/shimx64.efi']
+    return _find_efi_loader(loader_paths)
+
+
 def copy_windows_efi_bootloader(disk, partition):
     device = get_partition_device(disk, partition)
     mountpoint = device.replace('dev', 'mnt')
@@ -122,17 +144,8 @@ def copy_windows_efi_bootloader(disk, partition):
         umount(mountpoint)
         raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')
 
-    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/bootmgfw.efi',
-                    f'{esp_mountpoint}/EFI/Microsoft/Boot/bootmgfw.efi']
     try:
-        for efi_app in loader_paths:
-            logging.info(f'Searching for {efi_app}')
-            if os.path.exists(efi_app):
-                loader = efi_app
-                logging.info(f'Found bootloader at ESP partition: {loader}')
-                break
-        else:
-            raise OgError(f'Unable to locate Windows EFI bootloader bootmgfw.efi')
+        loader = find_windows_efi_loader(esp_mountpoint, bootlabel)
 
         loader_dir = os.path.dirname(loader)
         destination_dir = f'{mountpoint}/ogBoot'
-- 
cgit v1.2.3-18-g5258