diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 42930ba7..ba5eef3d 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -6,6 +6,7 @@ import logging import os import pathlib import platform +import re from contextlib import suppress import psutil @@ -71,6 +72,9 @@ KNOWN_HIVE_NAMES = { RAM_OK = 5.5 * 1024**3 # ~6 GiB assuming a bit of shared memory RAM_WARNING = 3.5 * 1024**3 # ~4 GiB assuming a bit of shared memory REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer' +REGEX_4K_ALIGNMENT = re.compile( + r'^(?P.*?)\s+(?P\d+)\s+(?P\d+)', + ) SLMGR = pathlib.Path(f'{os.environ.get("SYSTEMROOT")}/System32/slmgr.vbs') SYSTEMDRIVE = os.environ.get('SYSTEMDRIVE') @@ -163,19 +167,41 @@ def set_timezone(zone): # Info Functions def check_4k_alignment(show_alert=False): """Check if all partitions are 4K aligned, returns book.""" - cmd = ['WMIC', 'partition', 'get', 'StartingOffset'] + cmd = ['WMIC', 'partition', 'get', 'Caption,Size,StartingOffset'] + report = [] + show_alert = False # Check offsets proc = run_program(cmd) - for offset in proc.stdout.splitlines(): - offset = offset.strip() - if not offset.isnumeric(): + for line in proc.stdout.splitlines(): + line = line.strip() + if not line or not line.startswith('Disk'): continue - if int(offset) % 4096 != 0: - # Not aligned - if show_alert: - show_alert_box('One or more partitions are not 4K aligned') - raise GenericError('One or more partitions are not 4K aligned') + match = REGEX_4K_ALIGNMENT.match(line) + if not match: + LOG.error('Failed to parse partition info for: %s', line) + continue + if int(match.group('offset')) % 4096 == 0: + report.append( + color_string( + f'{match.group("description")}' + f' ({bytes_to_string(match.group("size"), decimals=1)})' + , + 'RED' + ) + ) + + # Show alert + if show_alert: + show_alert_box('One or more partitions not 4K aligned') + + # Done + if report: + report.insert( + 0, + color_string('One or more partitions not 4K aligned', 'YELLOW'), + ) + return report def export_bitlocker_info():