diff options
| -rw-r--r-- | src/utils/disk.py | 43 | ||||
| -rw-r--r-- | src/utils/postinstall.py | 35 | ||||
| -rw-r--r-- | src/utils/winreg.py | 23 |
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] |