1 files changed, 23 insertions, 0 deletions
diff --git a/src/utils/winreg.py b/src/utils/winreg.py
index ed2d5c6..0b22097 100644
--- a/src/utils/winreg.py
+++ b/src/utils/winreg.py
@@ -11,6 +11,8 @@ import hivex
 from enum import Enum
 from src.log import OgError
 from uuid import UUID
+from src.utils.disk import *
+from src.utils.uefi import is_uefi_supported
 
 
 WINDOWS_HIVES_PATH = '/Windows/System32/config'
@@ -112,3 +114,24 @@ def uuid_to_bytes(uuid):
     group4 = uuid[20:32]
     res = f'{group0}-{group1}-{group2}-{group3}-{group4}'
     return UUID(res).bytes
+
+
+def get_disk_id_bytes(disk):
+    disk_id = get_disk_id(disk)
+
+    if is_uefi_supported():
+        return uuid_to_bytes(disk_id)
+
+    return bytes.fromhex(disk_id)[::-1]
+
+
+def get_part_id_bytes(disk, partition):
+    if is_uefi_supported():
+        part_id = get_partition_id(disk, partition)
+        return uuid_to_bytes(part_id)
+
+    partition_start_offset = get_partition_start_offset(disk, partition)
+    sector_size = get_sector_size(disk)
+    byte_offset = partition_start_offset * sector_size
+    byte_offset = "{0:016x}".format(byte_offset)
+    return bytes.fromhex(byte_offset)[::-1]