summaryrefslogtreecommitdiffstats
path: root/src/utils
diff options
context:
space:
mode:
Diffstat (limited to 'src/utils')
-rw-r--r--src/utils/disk.py43
-rw-r--r--src/utils/postinstall.py35
-rw-r--r--src/utils/winreg.py23
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]