Added Windows Registry functions

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.")