disk: add get_efi_partition function

Add utility function inside disk.py to find, if any, the first ESP
partition of a given disk.

The disk is provided as an integer (starting at 1 following OpenGnsys
scripts usual values), meaning the (n-1)th disk from the disk array
returned from get_disks(). In the future a better mechanism should be
put in place to fetch probed disks from a running client.

This change is part of the upcoming drop of "IniciarSesion" script in
favor of a Python native approach. Specifically regarding UEFI systems.

1 files changed, 24 insertions, 0 deletions

diff --git a/src/utils/disk.py b/src/utils/disk.py
index 1f1fdab..02c63fc 100644
--- a/src/utils/disk.py
+++ b/src/utils/disk.py
@@ -34,3 +34,27 @@ def get_partition_device(disknum, partnum):
             return cxt.partition_to_string(pa, fdisk.FDISK_FIELD_DEVICE)
 
     raise ValueError(f'No such partition with disk index {disknum} and partition index {partnum}')
+
+
+def get_efi_partition(disknum):
+    """
+    Look for an EFI System Partition at the n-th disk. If disknum is invalid an exception is thrown
+
+    Returns tuple with:
+        - Device name containing the ESP
+        - /dev/{device} string
+        - Partition number (starting at 1)
+    """
+    try:
+        disk = get_disks()[disknum-1]
+        cxt = fdisk.Context(f'/dev/{disk}')
+
+        if cxt.label == fdisk.FDISK_DISKLABEL_DOS:
+            raise RuntimeError('Disk has DOS partition scheme, cannot find ESP.')
+
+        for pa in cxt.partitions:
+            if pa.type.name == 'EFI System':
+                return cxt.partition_to_string(pa, fdisk.FDISK_FIELD_DEVICE), f'/dev/{disk}', pa.partno + 1
+    except:
+        logging.error(f'Unable to find efi partition at disk number {disknum}')
+        raise