2 files changed, 44 insertions, 3 deletions
diff --git a/src/utils/disk.py b/src/utils/disk.py
index c424c61..94bbc65 100644
--- a/src/utils/disk.py
+++ b/src/utils/disk.py
@@ -16,6 +16,47 @@ from src.log import OgError
 import fdisk
 
 
+class ogPartitionData:
+    def __init__(self, parttype, fstype, padev, size, partno):
+        self.parttype = parttype
+        self.fstype = fstype
+        self.padev = padev
+        self.size = size
+        self.partno = partno
+
+class ogDiskData:
+    def __init__(self, nsectors, sector_size, label_name):
+        self.nsectors = nsectors
+        self.sector_size = sector_size
+        self.label_name = label_name
+
+def get_partition_data(device):
+    res = []
+    try:
+        cxt = fdisk.Context(device=device, details=True)
+    except Exception as e:
+        raise OgError(f'Partition query error: {e}') from e
+
+    for i, p in enumerate(cxt.partitions):
+        pd = ogPartitionData(
+            parttype = cxt.partition_to_string(p, fdisk.FDISK_FIELD_TYPEID),
+            fstype = cxt.partition_to_string(p, fdisk.FDISK_FIELD_FSTYPE),
+            padev = cxt.partition_to_string(p, fdisk.FDISK_FIELD_DEVICE),
+            size = cxt.partition_to_string(p, fdisk.FDISK_FIELD_SIZE),
+            partno = p.partno)
+        res.append(pd)
+    return res
+
+def get_disk_data(device):
+    try:
+        cxt = fdisk.Context(device=device, details=True)
+    except Exception as e:
+        raise OgError(f'Partition query error: {e}') from e
+    return ogDiskData(
+            nsectors = cxt.nsectors,
+            sector_size = cxt.sector_size,
+            label_name = cxt.label.name if cxt.label else "")
+
 def get_disks():
     """
     Walks /sys/block/ and returns files starting with 'sd',
diff --git a/src/utils/fstab.py b/src/utils/fstab.py
index b419aa2..05b94bc 100644
--- a/src/utils/fstab.py
+++ b/src/utils/fstab.py
@@ -158,11 +158,11 @@ def configure_swap(disk, mountpoint, fstab):
             swap_entry = entry
 
     diskname = get_disks()[disk-1]
-    cxt = fdisk.Context(f'/dev/{diskname}')
+    partitions = get_partition_data(device=f'/dev/{diskname}')
 
     swap_device = ''
-    for pa in cxt.partitions:
-        if cxt.partition_to_string(pa, fdisk.FDISK_FIELD_FSTYPE) == 'swap':
+    for pa in partitions:
+        if pa.fstype == 'swap':
             swap_device = get_formatted_device(disk, pa.partno + 1)
             break