diff --git a/.github/workflows/python-static-analysis-and-test.yml b/.github/workflows/python-static-analysis-and-test.yml index 4648b7d..9fccccb 100644 --- a/.github/workflows/python-static-analysis-and-test.yml +++ b/.github/workflows/python-static-analysis-and-test.yml @@ -49,6 +49,9 @@ jobs: python: ['3.7', '3.8', '3.9', '3.10'] runs-on: 'windows-latest' + # Enable registry write testing + env: + CASEMENT_TEST_WRITE_ENV: 1 steps: - name: Checkout code diff --git a/README.md b/README.md index 50b2bce..682d063 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,8 @@ A Python library that provides useful functionality for managing Microsoft Windo * Finding, creating and moving shortcuts. * Pinning shortcuts to the taskbar and start menu. * Flashing a window to get a users attention. +* Windows Registry reading and modification. +* Windows environment variable interface via registry allowing for permanent modification. ## Shortcut management @@ -43,6 +45,54 @@ Pinning/unpinning a shortcut to the taskbar: C:\blur\dev\casement>casement shortcut unpin "C:\Users\Public\Desktop\My Shortcut.lnk" -t ``` +# Environment Variables + +Access and modify environment variables stored in the registry using +`casement.env_var.EnvVar`. This lets you access the raw un-expanded file registry +values unlike what you see with os.environ. You can permanently modify the +environment variables for the system. +```python +from casement.env_var import EnvVar + +user_env = EnvVar(system=False) +user_env['VARIABLE_A'] = 'C:\\PROGRA~1' +expanded = user_env.normalize_path(user_env['VARIABLE_A'], tilde=False) +assert expanded == 'C:\\Program Files' +if 'VARIABLE_A' in user_env: + del user_env['VARIABLE_A'] +``` + +# Windows Registry + +Read and modify registry keys with `casement.registry.RegKey`. +```py +from casement.registry import RegKey + +reg = RegKey('HKLM', 'Software\\Autodesk\\Maya') +for ver in reg.child_names(): + # Attempt to access a specific sub-key that exists if ver is an version key + child = reg.child('{}\\Setup\\InstallPath'.format(ver)) + if child.exists(): + print(ver, child.entry('MAYA_INSTALL_LOCATION').value()) +``` + +There is a map of common registry locations like environment variables and classes. + +```python +from casement.registry import RegKey, REG_LOCATIONS + +# Access to the user and system keys driving `HKEY_CLASSES_ROOT` +classes_system = RegKey(*REG_LOCATIONS['system']['classes']) +classes_user = RegKey(*REG_LOCATIONS['user']['classes']) + +# Access to the user and system key controlling persistent Environment Variables +env_var_system = RegKey(*REG_LOCATIONS['system']['env_var']) +env_var_user = RegKey(*REG_LOCATIONS['user']['env_var']) + +# Uninstall key for "Add/Remove Programs" +uninstall_system = RegKey(*REG_LOCATIONS['system']['uninstall']) +``` + # Installing Casement can be installed by the standard pip command `pip install casement`. diff --git a/casement/env_var.py b/casement/env_var.py new file mode 100644 index 0000000..7c04f17 --- /dev/null +++ b/casement/env_var.py @@ -0,0 +1,154 @@ +import logging +import os +from contextlib import contextmanager + +import win32api +import win32con +import win32file +import win32gui +from six.moves import winreg +from six.moves.collections_abc import MutableMapping + +from .registry import REG_LOCATIONS, RegKey + +_logger_modify = logging.getLogger(__name__ + ".modify") +_logger_broadcast = logging.getLogger(__name__ + ".broadcast") + + +class EnvVar(MutableMapping): + """A dictionary like mapping of environment variables as defined in the + registry. Stored environment variables are not expanded by default. You can + use `normalize_path` for various path manipulation like expanding variables + and dealing with short/long file paths. + + Arguments: + system (bool, optional): Controls if this instance modifies system or + user environment variables. + + Every time you set or delete a environment variable it will send the + `WM_SETTINGCHANGE` message notifying other applications that system-wide + settings have changed. + See: https://learn.microsoft.com/en-us/windows/win32/winmsg/wm-settingchange + + You can prevent this message from being sent every time you modify env vars + by using the `EnvVar.delayed_broadcast()` with context: + + user_env = EnvVar(system=False) + with EnvVar.delayed_broadcast(): + user_env['VAR_B'] = 'value b' + user_env['VAR_C'] = 'value c' + del user_env['VAR_A'] + + By using `EnvVar.delayed_broadcast()` broadcast will only be called once, + instead of three times. + """ + + # These class variables are used to control if broadcast sends its message + _broadcast_enabled = True + _broadcast_required = False + + def __init__(self, system=False): + key, sub_key = REG_LOCATIONS["system" if system else "user"]['env_var'] + self.__reg__ = RegKey(key, sub_key) + + def __delitem__(self, key): + entry = self.__reg__.entry(key) + with self.delayed_broadcast(): + _logger_modify.debug('Deleting env var: "{}"'.format(key)) + entry.delete() + type(self)._broadcast_required = True + + def __getitem__(self, key): + entry = self.__reg__.entry(key) + try: + return entry.value() + except OSError: + raise KeyError(key) + + def __iter__(self): + return iter(self.__reg__.entry_names()) + + def __len__(self): + reg_key = self.__reg__._key() + if not reg_key: + return 0 + + _, count, _ = winreg.QueryInfoKey(reg_key) + return count + + def __setitem__(self, key, value, value_type=winreg.REG_EXPAND_SZ): + entry = self.__reg__.entry(key) + with self.delayed_broadcast(): + _logger_modify.debug('Setting env var: "{}" to "{}"'.format(key, value)) + entry.set(value, value_type=value_type) + type(self)._broadcast_required = True + + @classmethod + def broadcast(cls): + """Notify the system about the changes. Uses SendMessageTimeout to avoid + hanging if applications fail to respond to the broadcast. + """ + if cls._broadcast_enabled is False: + _logger_broadcast.debug( + 'Skipping broadcasting that the environment was changed' + ) + return + + _logger_broadcast.debug('Broadcasting that the environment was changed') + win32gui.SendMessageTimeout( + win32con.HWND_BROADCAST, + win32con.WM_SETTINGCHANGE, + 0, + 'Environment', + win32con.SMTO_ABORTIFHUNG, + 1000, + ) + # Reset the needs broadcast variable now that we have broadcast + cls._broadcast_required = False + + @classmethod + @contextmanager + def delayed_broadcast(cls): + """Context manager that disables broadcast calls for each set/delete of + environment variables. Only if any set/delete calls were made inside the + with context, will it call broadcast. + """ + try: + current = cls._broadcast_enabled + cls._broadcast_enabled = False + yield + finally: + cls._broadcast_enabled = current + if cls._broadcast_required: + cls.broadcast() + + @classmethod + def normalize_path(cls, path, expandvars=True, tilde=None, normpath=True): + """Normalize a path handling short/long names, env vars and normpath. + + Args: + path (str): Input path to expand. + expandvars (bool, optional): Expand environment variables. + tilde (bool, optional): Converts windows short paths(`C:\\PROGRA~1`). + If True, short paths are expanded to full file paths. If False + file paths are converted to short paths. Nothing is done if None. + normpath (bool, optional): If True calls os.path.normpath. + + Returns: + str: Input path with any tilde-shortenings expanded. + + Raises: + pywintypes.error: Raised if path doesn't exist when tilde is not None. + """ + if tilde: + path = win32api.GetShortPathName(path) + elif tilde is False: + path = win32file.GetLongPathName(path) + + if expandvars: + path = os.path.expandvars(path) + + if normpath: + path = os.path.normpath(path) + + return path diff --git a/casement/registry.py b/casement/registry.py new file mode 100644 index 0000000..1210c10 --- /dev/null +++ b/casement/registry.py @@ -0,0 +1,218 @@ +""" +key: Nodes in tree +Each node in the tree is called a key. +Each key can contain both subkeys and data entries called values. +""" +import six +from six.moves import winreg + + +def _generate_key_map(): + """Creates a map of strings to winreg.HKEY_* objects.""" + key_map = {} + + for key in dir(winreg): + if key.startswith('HKEY_'): + # Add the long names + key_map[key] = getattr(winreg, key) + # Add the winreg value so we can lookup nice names for keys + key_map[getattr(winreg, key)] = key + # Generate shortnames + short = 'HK' + ''.join([s[0] for s in key.split('_')[1:]]) + if short not in key_map: + key_map[short] = key_map[key] + return key_map + + +key_map = _generate_key_map() + +"""Stores commonly looked up locations in the registry.""" +REG_LOCATIONS = { + "system": { + "classes": (winreg.HKEY_LOCAL_MACHINE, 'Software\\Classes'), + "env_var": ( + winreg.HKEY_LOCAL_MACHINE, + 'SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment', + ), + "uninstall": ( + winreg.HKEY_LOCAL_MACHINE, + 'Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall', + ), + }, + "user": { + "classes": (winreg.HKEY_CURRENT_USER, 'Software\\Classes'), + "env_var": (winreg.HKEY_CURRENT_USER, 'Environment'), + }, +} + + +class RegKey(object): + def __init__(self, key, sub_key=None, computer_name=None, architecture=64): + # If sub_key is not passed, it must be part of key + if sub_key is None: + key, sub_key = key.split('\\', 1) + + # Convert key strings into winreg keys + if isinstance(key, six.string_types): + key = key.upper() + key = key_map[key] + + self.architecture = architecture + self.computer_name = computer_name + self.key = key + self.sub_key = sub_key + + def __str__(self): + return 'RegKey({!r}, {!r}, {!r}, {})'.format( + key_map.get(self.key, self.key), + self.sub_key, + self.computer_name, + self.architecture, + ) + + def _key(self, write=False, create=False, delete=False): + """Returns a winreg.OpenKeyEx or None. + + Returns: + A winreg handle object + """ + connection = winreg.ConnectRegistry(None, self.key) + sam = self._sam(self.architecture) + access = winreg.KEY_READ + if write: + access = winreg.KEY_WRITE + + if create: + func = winreg.CreateKeyEx + else: + func = winreg.OpenKeyEx + + try: + return func(connection, self.sub_key, 0, access | sam) + except OSError: + return None + + @classmethod + def _sam(cls, architecture): + """Return the correct wow64 key for the requested architecture.""" + if architecture == 32: + return winreg.KEY_WOW64_32KEY + elif architecture == 64: + return winreg.KEY_WOW64_64KEY + return 0 + + def child(self, name): + """Returns a RegKey for a sub_key with the given name. + + Args: + name (str): The name of the child key to return. + """ + cls = type(self) + sub_key = '{}\\{}'.format(self.sub_key, name) + return cls( + self.key, + sub_key=sub_key, + computer_name=self.computer_name, + architecture=self.architecture, + ) + + def child_names(self): + """Generator returning the name for all sub_keys of this key.""" + reg_key = self._key() + if not reg_key: + return + + index = 0 + while True: + try: + yield winreg.EnumKey(reg_key, index) + index += 1 + except OSError: + break + + def create(self): + self._key(create=True) + + def delete(self): + """Delete the entire registry key. Raises a RuntimeError if the key has + child keys.""" + # If the key doesn't exist, there is nothing to do + if not self.exists(): + return False + + # TODO: add recursive delete? + for child in self.child_names(): + # We can't delete this key if it has sub-keys. Using for loop instead + # of using next() so this can be used inside a contextmanager + if child: + msg = "Unable to delete key, it has sub-keys. {}".format(self) + raise RuntimeError(msg) + + sam = self._sam(self.architecture) + winreg.DeleteKeyEx(self.key, self.sub_key, sam) + return True + + def entry(self, name=None): + """Returns a RegEntry for the given name. To access the (Default) entry + pass an empty string or None. + + Args: + name (str, optional): The name of the entry to return. + """ + return RegEntry(self, name) + + def entry_names(self): + """Generator returning the name for all entry's stored on this key.""" + reg_key = self._key() + if not reg_key: + return + + _, count, _ = winreg.QueryInfoKey(reg_key) + for index in range(count): + name, _, _ = winreg.EnumValue(reg_key, index) + yield name + + def exists(self): + """Returns True if the key exists in the registry.""" + key = self._key() + return key is not None + + +class RegEntry(object): + def __init__(self, key, name): + self.key = key + self.name = name + + def delete(self): + key = self.key._key(write=True) + winreg.DeleteValue(key, self.name) + winreg.CloseKey(key) + + def type(self): + """Returns the winreg value type.""" + return self.value_info()[1] + + def set(self, value, value_type=None): + key = self.key._key(write=True, create=True) + + if isinstance(value_type, six.string_types): + value_type = getattr(winreg, value_type) + + winreg.SetValueEx(key, self.name, 0, value_type, value) + winreg.CloseKey(key) + # TODO: add notify, or keep that in EnvVar? + + def value(self): + """Returns the data value.""" + value, _ = self.value_info() + # We may need to do some conversion based on type in the future. + return value + + def value_info(self): + """Returns the entry value and value type. + + Returns: + object: Value stored in key + int: registry type for value. See winreg's Value Types + """ + return winreg.QueryValueEx(self.key._key(), self.name) diff --git a/setup.cfg b/setup.cfg index f0a374d..ef0c19d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -70,6 +70,7 @@ extend-ignore = E501, E722, W503, + B904, max-line-length = 80 exclude = *.egg-info diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..d134c0d --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,12 @@ +import os + +# Tests that write to the registry are disabled by default. Any test that is +# writing to the registry should use `pytest.mark.skipif` checking the env var +# `CASEMENT_TEST_WRITE_ENV`. It's intended that this is only enabled on the +# github action runner, and by developers that need to modify code modifying the +# registry and EnvVar setting/deleting features. +ENABLE_ENV_VAR = 'CASEMENT_TEST_WRITE_ENV' +SKIP_ENV_VAR_WRITES = os.getenv(ENABLE_ENV_VAR) != '1' +ENV_VAR_REASON = "To enable registry write tests, set the env var `{}` to `1`.".format( + ENABLE_ENV_VAR +) diff --git a/tests/test_env_var.py b/tests/test_env_var.py new file mode 100644 index 0000000..eaeb94e --- /dev/null +++ b/tests/test_env_var.py @@ -0,0 +1,209 @@ +""" +By default all registry tests that actually modify the registry are skipped. +You must set the environment variable `CASEMENT_TEST_WRITE_ENV` to `1` to enable +running the skipped tests. + +Environment variable modification tests should only affect casement specific +environment variables that don't affect anything. They should contain +`CASEMENT_DELETE_ME` in the name. When testing finishes all of these env var's +should be removed. All write tests should be done on user environment variables. + +Here are the Environment variables that are modified when write testing is enabled. +- `CASEMENT_TEST_DELETE_ME_ENV_VAR` + +Note: See test_registry.py to see how testing registry keys/entries are handled. +""" +from __future__ import absolute_import + +import os + +import pytest +import win32api + +from casement.env_var import EnvVar +from casement.registry import RegKey + +from . import ENV_VAR_REASON, SKIP_ENV_VAR_WRITES + +TEST_VAR_NAME = 'CASEMENT_TEST_DELETE_ME_ENV_VAR' + + +def env_var(var_name): + def decorator(function): + def wrapper(*args, **kwargs): + senv = EnvVar(True) + uenv = EnvVar(False) + # Ensure the test env var is cleaned up, in case a previous test failed + # or was killed before the env var was removed normally. + if var_name in senv: + del senv[var_name] + if var_name in uenv: + del uenv[var_name] + + ret = function(*args, **kwargs) + + # Ensure the variables are cleaned up after the test runs + if var_name in senv: + del senv[var_name] + if var_name in uenv: + del uenv[var_name] + return ret + + return wrapper + + return decorator + + +def test_expand_path(): + # Generate a short name path. This method needs a file that actually + # exists on disk, and test_env_var.py is long enough to get ~'ed + short_name = win32api.GetShortPathName(__file__) + assert '~' in short_name + + # Check expandvars argument + out = EnvVar.normalize_path('%HOMEPATH%', expandvars=True) + assert out == os.path.expandvars('%HOMEPATH%') + out = EnvVar.normalize_path('%HOMEPATH%', expandvars=False) + assert out == '%HOMEPATH%' + + # Check tilde argument + out = EnvVar.normalize_path(short_name, tilde=False) + assert os.path.normpath(out) == __file__ + out = EnvVar.normalize_path(__file__, tilde=True) + assert os.path.normpath(out) == os.path.normpath(short_name) + + # Check normpath argument + fle = __file__.replace('\\', '/') + out = EnvVar.normalize_path(fle, normpath=False) + assert out == fle + out = EnvVar.normalize_path(fle, normpath=True) + assert out == __file__ + + +@pytest.mark.skipif(SKIP_ENV_VAR_WRITES, reason=ENV_VAR_REASON) +@env_var(TEST_VAR_NAME) +def test_envvar_modify(): + uenv = EnvVar(system=False) + # Ensure the env var is not set + assert TEST_VAR_NAME not in uenv + + # Ensure the env var was set + uenv[TEST_VAR_NAME] = 'TEST' + assert uenv[TEST_VAR_NAME] == 'TEST' + uenv[TEST_VAR_NAME] = 'TEST 2' + assert uenv[TEST_VAR_NAME] == 'TEST 2' + + # Test iterating over the env var. We can't tell for sure what + # env vars are set, but we at least know one is set at this point + for key, value in uenv.items(): + if key == TEST_VAR_NAME: + assert value == 'TEST 2' + break + else: + raise AssertionError( + '{} was not found when iterating over the user environment'.format( + TEST_VAR_NAME + ) + ) + + # Test length support, again we only know at least one env var is set + count = len(uenv) + assert count >= 1 + + # Ensure the env var is once again not set + del uenv[TEST_VAR_NAME] + assert TEST_VAR_NAME not in uenv + + # test edge case where the registry key doesn't exist + bad = EnvVar(system=False) + bad.__reg__ = RegKey('HKEY_CURRENT_USER', 'BadEnvironment') + assert len(bad) == 0 + + +class EnvVarPatch(EnvVar): + """Replaces the broadcast classmethod with one that allows us to track + every time its called.""" + + count = 0 + + @classmethod + def broadcast(cls): + """Enable us to check how many times broadcast was called.""" + if cls._broadcast_enabled is False: + return + cls._broadcast_required = False + cls.count += 1 + + +@pytest.mark.skipif(SKIP_ENV_VAR_WRITES, reason=ENV_VAR_REASON) +@env_var(TEST_VAR_NAME) +def test_broadcast(): + """This test tests that delayed_broadcast only calls broadcast in the outer + most use of its with context, and only if a broadcast was required. + """ + uenv = EnvVarPatch(system=False) + + # Test setting an env var broadcasts + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 0 + uenv[TEST_VAR_NAME] = 'TEST' + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 1 + + # Test deleting an env var broadcasts + EnvVarPatch.count = 0 + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + del uenv[TEST_VAR_NAME] + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 1 + + +@pytest.mark.skipif(SKIP_ENV_VAR_WRITES, reason=ENV_VAR_REASON) +@env_var(TEST_VAR_NAME) +def test_delayed_broadcast(): + """This test tests that delayed_broadcast only calls broadcast in the outer + most use of its with context, and only if a broadcast was required. + """ + uenv = EnvVarPatch(system=False) + + # Test that multiple set/delete calls only emit a single broadcast + # when the outer with context exits. + EnvVarPatch.count = 0 + with EnvVarPatch.delayed_broadcast(): + assert EnvVarPatch._broadcast_enabled is False + assert EnvVarPatch._broadcast_required is False + uenv[TEST_VAR_NAME] = 'TEST' + assert EnvVarPatch._broadcast_enabled is False + assert EnvVarPatch._broadcast_required is True + assert EnvVarPatch.count == 0 + + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 1 + + # Test that no broadcast is made if no changes were made + EnvVarPatch.count = 0 + with EnvVarPatch.delayed_broadcast(): + assert EnvVarPatch._broadcast_enabled is False + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 0 + + assert EnvVarPatch._broadcast_enabled is True + assert EnvVarPatch._broadcast_required is False + assert EnvVarPatch.count == 0 + + # Test that the real broadcast method respects _broadcast_enabled + uenv = EnvVar(system=False) + with EnvVar.delayed_broadcast(): + assert EnvVar._broadcast_enabled is False + assert EnvVar._broadcast_required is False + uenv[TEST_VAR_NAME] = 'TEST' + assert EnvVar._broadcast_enabled is False + assert EnvVar._broadcast_required is True + + assert EnvVar._broadcast_enabled is True + assert EnvVar._broadcast_required is False diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000..0fecb52 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,222 @@ +""" +By default all registry tests that actually modify the registry are skipped. +You must set the environment variable `CASEMENT_TEST_WRITE_ENV` to `1` to enable +running the skipped tests. + +Registry modification tests should only affect casement specific registry keys +that don't affect anything. They should contain `CASEMENT_DELETE_ME` in the name. +When testing finishes all of these registry keys should be removed. All write +tests should be done on user hives not system ones. + +Here are the registry keys that are modified when write testing is enabled. +- `HKEY_CURRENT_USER\\Software\\Classes\\CASEMENT_DELETE_ME` and children + +Note: See test_env_var.py to see how testing Environment variables are handled. +""" +# TODO: Look into using a custom testing registry hive to handle all testing +# without the need for the host to actually have these registry keys. We should +# only enable testing of registry modifications once this is resolved. +# Ie don't running casement tests should not modify the host registry. +from __future__ import absolute_import + +from contextlib import contextmanager + +import pytest +import six +from six.moves import winreg + +from casement.registry import REG_LOCATIONS, RegKey + +from . import ENABLE_ENV_VAR, ENV_VAR_REASON, SKIP_ENV_VAR_WRITES + + +@contextmanager +def remove_reg_keys(items): + """Ensures the requested registry keys are removed before and after the + with context code is run. This ensures a consistent testing environment + even if a previous test was killed without cleaning up after itself. + + Items is a list of registry (key, sub_key) items. If removing parents and + children ensure the children are passed after the parents. + """ + items = reversed(items) + for key, sub_key in items: + reg = RegKey(key, sub_key) + # Ensure the registry key is removed, in case a previous test failed + # or was killed before the registry key was removed normally. + if reg.exists(): + reg.delete() + + try: + yield + finally: + for key, sub_key in items: + reg = RegKey(key, sub_key) + # Ensure the variables are cleaned up after the test runs + if reg.exists(): + reg.delete() + + +@pytest.mark.parametrize( + 'key,sub_key,hkey', + [ + # Note: Lowercase testing is handled in the test, only use uppercase here + ('HKEY_LOCAL_MACHINE', 'SOFTWARE\\Microsoft', winreg.HKEY_LOCAL_MACHINE), + ('HKLM', 'SOFTWARE\\Microsoft', winreg.HKEY_LOCAL_MACHINE), + ('HKEY_CURRENT_USER', 'SOFTWARE\\Microsoft', winreg.HKEY_CURRENT_USER), + ('HKCU', 'SOFTWARE\\Microsoft', winreg.HKEY_CURRENT_USER), + ('HKEY_CLASSES_ROOT', '*', winreg.HKEY_CLASSES_ROOT), + ('HKCR', 'Directory', winreg.HKEY_CLASSES_ROOT), + ('HKEY_USERS', '.DEFAULT\\Software\\Microsoft', winreg.HKEY_USERS), + ('HKU', '.DEFAULT\\Software\\Microsoft', winreg.HKEY_USERS), + ], +) +def test_reg_open(key, sub_key, hkey): + # TEST passing split key an sub_key's + # key and sub_key are passed in as 2 arguments + regkey = RegKey(key, sub_key) + assert regkey.key == hkey + assert regkey.sub_key == sub_key + # Key is lower cased + regkey = RegKey(key.lower(), sub_key) + assert regkey.key == hkey + assert regkey.sub_key == sub_key + + # Test passing just key argument with sub_key attached + regkey = RegKey('{}\\{}'.format(key, sub_key)) + assert regkey.key == hkey + assert regkey.sub_key == sub_key + + +def test_child_names(): + # Test valid registry key + reg = RegKey('HKLM\\SOFTWARE\\Microsoft\\Windows') + names = list(reg.child_names()) + assert 'Notepad' in names + assert 'Shell' in names + + # Test invalid registry key + reg = RegKey('HKLM\\SOFTWARE\\InvalidName') + assert list(reg.child_names()) == [] + + +def test_child(): + sub_key = 'SOFTWARE\\Microsoft\\Windows' + reg = RegKey('HKLM', sub_key) + notepad = reg.child('Notepad') + assert notepad._key() is not None + assert notepad.sub_key == '{}\\Notepad'.format(sub_key) + shell = reg.child('Shell') + assert shell._key() is not None + assert shell.sub_key == '{}\\Shell'.format(sub_key) + + +def test_entry_names(): + # Test Valid registry key + reg = RegKey('HKCR\\Directory') + names = reg.entry_names() + assert 'AlwaysShowExt' in names + + # Test invalid registry key + reg = RegKey('HKCR\\InvalidName') + assert list(reg.entry_names()) == [] + + +def test_entry(): + # Note: This test may break if the host computer has modified Directory + # settings. Just ensure this test passes when run by github actions. + reg = RegKey('HKCR\\Directory') + # `(Default)` entry + entry = reg.entry() + txt = 'File Folder' + assert entry.type() == winreg.REG_SZ + assert entry.value() == txt + assert entry.value_info() == (txt, winreg.REG_SZ) + + # Binary value type + entry = reg.entry('EditFlags') + assert entry.type() == winreg.REG_BINARY + assert entry.value() == b'\xd2\x01\x00\x00' + + # String data type + entry = reg.entry('PreviewTitle') + assert entry.type() == winreg.REG_SZ + assert isinstance(entry.value(), six.string_types) + assert entry.value() == 'prop:System.ItemNameDisplay;System.ItemTypeText' + + # Expanding string data type + entry = reg.entry('FriendlyTypeName') + assert entry.type() == winreg.REG_EXPAND_SZ + assert isinstance(entry.value(), six.string_types) + + +def test_sam(): + assert RegKey._sam(32) == winreg.KEY_WOW64_32KEY + assert RegKey._sam(64) == winreg.KEY_WOW64_64KEY + assert RegKey._sam(0) == 0 + assert RegKey._sam(16) == 0 + + +@pytest.mark.skipif(SKIP_ENV_VAR_WRITES, reason=ENV_VAR_REASON) +def test_write(): + key, root = REG_LOCATIONS['user']['classes'] + root = '{}\\CASEMENT_DELETE_ME'.format(root) + child_names = ('CASEMENT_DELETE_ME_CHILD1', 'CASEMENT_DELETE_ME_CHILD2') + child1 = '{}\\{}'.format(root, child_names[0]) + child2 = '{}\\{}'.format(root, child_names[1]) + + with remove_reg_keys(((key, root), (key, child1), (key, child2))): + reg = RegKey(key, root) + assert not reg.exists() + + # Test deleting a key that doesn't exist + ret = reg.delete() + assert ret is False + + # Test creating a new key + reg.create() + assert reg.exists() + + # Ensure we raise a useful exception if trying to delete a key with a child + creg1 = RegKey(key, child1) + assert not creg1.exists() + creg1.create() + assert creg1.exists() + with pytest.raises(RuntimeError) as excinfo: + reg.delete() + assert str(excinfo.value).startswith("Unable to delete key") + + # Test child and child_names by adding a second child + creg2 = reg.child(child_names[1]) + assert not creg2.exists() + creg2.create() + + assert set(reg.child_names()) == set(child_names) + + # Test getting/setting an Entry value + entry_name = 'CASEMENT_DELETE_ME_ENTRY' + entry = creg1.entry(entry_name) + assert entry.key == creg1 + assert entry.name == entry_name + + value = '[%{}%]'.format(ENABLE_ENV_VAR) + entry.set(value, "REG_EXPAND_SZ") + assert entry.value() == value + assert entry.type() == winreg.REG_EXPAND_SZ + assert entry.value_info() == (value, winreg.REG_EXPAND_SZ) + entry.set(value, winreg.REG_SZ) + assert entry.value_info() == (value, winreg.REG_SZ) + + # Remove the children so we can actually delete the parent + creg1.delete() + assert not creg1.exists() + creg2.delete() + assert not creg2.exists() + + # Test deleting a key that exists + ret = reg.delete() + assert ret is True + assert not reg.exists() + + # Check child_names if the key doesn't exist + assert not set(reg.child_names()) diff --git a/tox.ini b/tox.ini index a997d94..d225aa1 100644 --- a/tox.ini +++ b/tox.ini @@ -5,6 +5,9 @@ skipsdist = True [testenv] changedir = {toxinidir} +# This env var is used to enable registry write tests by developers and the +# gitlab action runner. +passenv = CASEMENT_TEST_WRITE_ENV setenv = {py27}: PYTHONWARNINGS=ignore:DEPRECATION::pip._internal.cli.base_command skip_install = True