From 11a963c70930405cb94d222ec5a55b1e5000bcb4 Mon Sep 17 00:00:00 2001 From: Alejandro Sirgo Rica Date: Tue, 30 Apr 2024 10:26:44 +0200 Subject: 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. --- src/utils/disk.py | 43 +++++++++++++++++++++++++++++++++++++++++++ src/utils/postinstall.py | 35 ++++++++++++++++++++++++++--------- src/utils/winreg.py | 23 +++++++++++++++++++++++ 3 files changed, 92 insertions(+), 9 deletions(-) (limited to 'src') 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] -- cgit v1.2.3-18-g5258