diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 7f17af37..77b61999 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -5,7 +5,9 @@ import logging import os import pathlib import platform +import winreg +from contextlib import suppress from wk.borrowed import acpi from wk.exe import run_program from wk.io import non_clobber_path @@ -15,6 +17,39 @@ from wk.std import GenericError, GenericWarning, sleep # STATIC VARIABLES LOG = logging.getLogger(__name__) +KNOWN_HIVES = { + 'HKCR': winreg.HKEY_CLASSES_ROOT, + 'HKCU': winreg.HKEY_CURRENT_USER, + 'HKLM': winreg.HKEY_LOCAL_MACHINE, + 'HKU': winreg.HKEY_USERS, + 'HKEY_CLASSES_ROOT': winreg.HKEY_CLASSES_ROOT, + 'HKEY_CURRENT_USER': winreg.HKEY_CURRENT_USER, + 'HKEY_LOCAL_MACHINE': winreg.HKEY_LOCAL_MACHINE, + 'HKEY_USERS': winreg.HKEY_USERS, + } +KNOWN_HIVE_NAMES = { + winreg.HKEY_CLASSES_ROOT: 'HKCR', + winreg.HKEY_CURRENT_USER: 'HKCU', + winreg.HKEY_LOCAL_MACHINE: 'HKLM', + winreg.HKEY_USERS: 'HKU', + winreg.HKEY_CLASSES_ROOT: 'HKEY_CLASSES_ROOT', + winreg.HKEY_CURRENT_USER: 'HKEY_CURRENT_USER', + winreg.HKEY_LOCAL_MACHINE: 'HKEY_LOCAL_MACHINE', + winreg.HKEY_USERS: 'HKEY_USERS', + } +KNOWN_VALUE_TYPES = { + 'BINARY': winreg.REG_BINARY, + 'DWORD': winreg.REG_DWORD, + 'DWORD_LITTLE_ENDIAN': winreg.REG_DWORD_LITTLE_ENDIAN, + 'DWORD_BIG_ENDIAN': winreg.REG_DWORD_BIG_ENDIAN, + 'EXPAND_SZ': winreg.REG_EXPAND_SZ, + 'LINK': winreg.REG_LINK, + 'MULTI_SZ': winreg.REG_MULTI_SZ, + 'NONE': winreg.REG_NONE, + 'QWORD': winreg.REG_QWORD, + 'QWORD_LITTLE_ENDIAN': winreg.REG_QWORD_LITTLE_ENDIAN, + 'SZ': winreg.REG_SZ, + } OS_VERSION = float(platform.win32_ver()[0]) REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer' SLMGR = pathlib.Path(f'{os.environ.get("SYSTEMROOT")}/System32/slmgr.vbs') @@ -177,5 +212,153 @@ def run_sfc_scan(): raise OSError +# Registry Functions +def reg_delete_key(hive, key, recurse=False): + """Delete a key from the registry. + + NOTE: If recurse is False then it will only work on empty keys. + """ + hive = reg_get_hive(hive) + hive_name = KNOWN_HIVE_NAMES.get(hive, '???') + + # Delete subkeys first + if recurse: + with suppress(WindowsError), winreg.OpenKey(hive, key) as open_key: + while True: + subkey = fr'{key}\{winreg.EnumKey(open_key, 0)}' + reg_delete_key(hive, subkey, recurse=recurse) + + # Delete key + try: + winreg.DeleteKey(hive, key) + LOG.warning(r'Deleting registry key: %s\%s', hive_name, key) + except FileNotFoundError: + # Ignore + pass + except PermissionError: + LOG.error(r'Failed to delete registry key: %s\%s', hive_name, key) + if recurse: + # Re-raise exception + raise + + # recurse is not True so assuming we tried to remove a non-empty key + msg = fr'Refusing to remove non-empty key: {hive_name}\{key}' + raise FileExistsError(msg) + + +def reg_delete_value(hive, key, value): + """Delete a value from the registry.""" + access = winreg.KEY_ALL_ACCESS + hive = reg_get_hive(hive) + hive_name = KNOWN_HIVE_NAMES.get(hive, '???') + + # Delete value + with winreg.OpenKey(hive, key, access=access) as open_key: + try: + winreg.DeleteValue(open_key, value) + LOG.warning( + r'Deleting registry value: %s\%s "%s"', hive_name, key, value, + ) + except FileNotFoundError: + # Ignore + pass + except PermissionError: + LOG.error( + r'Failed to delete registry value: %s\%s "%s"', hive_name, key, value, + ) + # Re-raise exception + raise + + +def reg_get_hive(hive): + """Get winreg HKEY constant from string, returns HKEY constant.""" + if isinstance(hive, int): + # Assuming we're already a winreg HKEY constant + pass + else: + hive = KNOWN_HIVES[hive.upper()] + + # Done + return hive + + +def reg_get_value_type(value_type): + """Get registry value type from string, returns winreg constant.""" + if isinstance(value_type, int): + # Assuming we're already a winreg value type constant + pass + else: + value_type = KNOWN_VALUE_TYPES[value_type.upper()] + + # Done + return value_type + + +def reg_key_exists(hive, key): + """Test if the specified hive/key exists, returns bool.""" + exists = False + hive = reg_get_hive(hive) + + # Query key + try: + winreg.QueryValue(hive, key) + except FileNotFoundError: + # Leave set to False + pass + else: + exists = True + + # Done + return exists + + +def reg_read_value(hive, key, value, force_32=False, force_64=False): + """Query value from hive/hey, returns multiple types.""" + access = winreg.KEY_READ + data = None + hive = reg_get_hive(hive) + + # Set access + if force_32: + access = access | winreg.KEY_WOW64_32KEY + elif force_64: + access = access | winreg.KEY_WOW64_64KEY + + # Query value + with winreg.OpenKey(hive, key, access=access) as open_key: + # Returning first part of tuple and ignoreing type + data = winreg.QueryValueEx(open_key, value)[0] + + # Done + return data + + +def reg_write_settings(settings_dict): + """TODO""" + + +def reg_set_value( + hive, key, name, value, value_type, force_32=False, force_64=False): + # pylint: disable=too-many-arguments + """Set value for hive/key.""" + access = winreg.KEY_WRITE + value_type = reg_get_value_type(value_type) + hive = reg_get_hive(hive) + + # Set access + if force_32: + access = access | winreg.KEY_WOW64_32KEY + elif force_64: + access = access | winreg.KEY_WOW64_64KEY + + # Create key + winreg.CreateKeyEx(hive, key, access=access) + + # Set value + with winreg.OpenKey(hive, key, access=access) as open_key: + # Returning first part of tuple and ignoreing type + winreg.SetValueEx(open_key, name, 0, value_type, value) + + if __name__ == '__main__': print("This file is not meant to be called directly.")