utils: add disk index checks in get_partition_device

Control the possibility of a bad disk index. Handle the checks in
a similar fashion to the get_efi_partition function.

1 files changed, 5 insertions, 1 deletions

diff --git a/src/utils/disk.py b/src/utils/disk.py
index 0848cff..b4541cc 100644
--- a/src/utils/disk.py
+++ b/src/utils/disk.py
@@ -27,7 +27,11 @@ def get_partition_device(disknum, partnum):
     """
     Returns the device path, given a disk and partition number
     """
-    disk = get_disks()[disknum-1]
+    disk_index = disknum - 1
+    if disk_index < 0 or disk_index >= len(get_disks()):
+        raise ValueError(f'Invalid disk number {disknum}, {len(get_disks())} disks available.')
+
+    disk = get_disks()[disk_index]
     cxt = fdisk.Context(f'/dev/{disk}')
 
     for pa in cxt.partitions: