1 files changed, 41 insertions, 0 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',