utils: Add windows_register_c_drive implementation

Add python implementation of the legacy ogWindowsRegisterPartition
function.

This function configures the system partition to be considered the
new C drive after a Windows image restore.

The drive letter configuration is stored in the SYSTEM hive of
the windows registry. The node MountedDevices contains all the
configuration as key-value pairs.
The C drive key \DosDevices\C: contains a different value based on
the partitioning type.

GPT
the value is DMIO:ID: followed by the 16 byte GUID of the
partition.
Example: DMIO:ID:\xc9\xfc\x1c\x86\x13\x11O\t\xa2\x9c{/\xf1\xdf\xe4)

MBR
The value is a little endian byte sequence with 2 parts.
The first 4 bytes correspond to the disk id. The second part is
the byte offset start of the partition in that disk.
Example: \xe1\\\x9cP\x00\x00\x10j\x18\x00\x00\x00

If we format the MBR value in a more readable way we get
e1 5c 9c 50 00 00 10 6a 18 00 00 00
In this case the disk ID is 509c5ce1.
The partition offset is 186a100000.

This patch adds the following helper functions to:

- get_disk_id_bytes(): to obtain the disk identifier.

- get_part_id_bytes(): to obtain a partition identifier as UUID
  for MBR or DMIO:ID format for GPT.

- get_sector_size(): to query the sector size of a specific disk.
  Read /sys/class/block/{device_name}/queue/hw_sector_size to obtain the value.
  This is MBR specific.

- get_partition_start_offset(): to query the start sector of a specific
  partition and disk. Use sfdisk with the -J argument to get fdisk data in
  json format. This is MBR specific.

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]