3 files changed, 92 insertions, 9 deletions

diff --git a/src/utils/disk.py b/src/utils/disk.py
index d7706fe..c424c61 100644
--- a/src/utils/disk.py
+++ b/src/utils/disk.py
@@ -10,6 +10,7 @@ import os
 import logging
 import shlex
 import subprocess
+import json
 from src.log import OgError
 
 import fdisk
@@ -106,3 +107,45 @@ def get_filesystem_id(disk_index, part_index):
     if proc.returncode != 0:
             raise OgError(f'failed to query filesystem UUID for {device}')
     return proc.stdout.strip()
+
+
+def get_sector_size(disk):
+    disk_index = disk - 1
+
+    if disk_index < 0 or disk_index >= len(get_disks()):
+        raise OgError(f'Invalid disk number {disk} when trying to find ESP, {len(get_disks())} disks available.')
+
+    device_name = get_disks()[disk_index]
+    file_path = f'/sys/class/block/{device_name}/queue/hw_sector_size'
+
+    try:
+        with open(file_path, 'r') as f:
+            data = f.read().strip()
+    except OSError as e:
+        raise OgError(f'Error while trying to read {file_path}: {e}') from e
+    return int(data)
+
+
+def get_partition_start_offset(disk, partition):
+    disk_name = get_disks()[disk - 1]
+    disk_path = f'/dev/{disk_name}'
+    part_number = partition - 1
+
+    cmd = f'sfdisk -J {disk_path}'
+    proc = subprocess.run(shlex.split(cmd), capture_output=True, text=True)
+
+    if proc.returncode != 0:
+        raise OgError(f'Failed to query sfdisk')
+
+    try:
+        part_data_json = json.loads(proc.stdout)
+    except json.JSONDecodeError as e:
+        raise OgError(f'Invalid sfdisk output: {e}') from e
+
+    try:
+        part_data = part_data_json['partitiontable']['partitions']
+        start_offset = part_data[part_number]['start']
+    except KeyError as e:
+        raise OgError(f'Error while trying to parse sfdisk: {e}') from e
+
+    return start_offset
diff --git a/src/utils/postinstall.py b/src/utils/postinstall.py
index 94a7e4b..221da81 100644
--- a/src/utils/postinstall.py
+++ b/src/utils/postinstall.py
@@ -110,16 +110,33 @@ def configure_os_custom(disk, partition):
 
 
 def windows_register_c_drive(disk, partition):
-    cmd_configure = f"ogWindowsRegisterPartition {disk} {partition} C {disk} {partition}"
+    device = get_partition_device(disk, partition)
+    mountpoint = device.replace('dev', 'mnt')
 
-    proc = subprocess.run(cmd_configure,
-                          stdout=subprocess.DEVNULL,
-                          stderr=subprocess.DEVNULL,
-                          encoding='utf-8',
-                          shell=True,
-                          check=True)
-    if proc.returncode != 0:
-        logging.warning(f'{cmd_configure} returned non-zero exit status {proc.returncode}')
+    if not mount_mkdir(device, mountpoint):
+        raise OgError(f'Unable to mount {device} into {mountpoint}')
+
+    hive_path = f'{mountpoint}{WINDOWS_HIVE_SYSTEM}'
+
+    try:
+        hive = hive_handler_open(hive_path, write=True)
+        root = hive.root()
+        device_node = get_node_child_from_path(hive, root, 'MountedDevices')
+
+        if is_uefi_supported():
+            part_id = get_part_id_bytes(disk, partition)
+            value = b'DMIO:ID:' + part_id
+        else:
+            disk_id = get_disk_id_bytes(disk)
+            part_offset = get_part_id_bytes(disk, partition)
+
+            value = bytes(disk_id + part_offset)
+
+        device_value = {'key': '\DosDevices\C:', 't': RegistryType.BINARY.value, 'value': value}
+        hive.node_set_value(device_node, device_value)
+        hive.commit(hive_path)
+    finally:
+        umount(mountpoint)
 
 
 def configure_mbr_boot_sector(disk, partition):
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]