From ddf08779ae8db05a75d51e063f507d1c3bfa148e Mon Sep 17 00:00:00 2001 From: Alejandro Sirgo Rica Date: Mon, 11 Mar 2024 09:44:38 +0100 Subject: 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. --- src/utils/disk.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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: -- cgit v1.2.3-18-g5258