diff --git a/dfwinreg/interface.py b/dfwinreg/interface.py index 6e33c4d..1ab04d0 100644 --- a/dfwinreg/interface.py +++ b/dfwinreg/interface.py @@ -10,10 +10,6 @@ class WinRegistryFile(object): """Windows Registry file interface.""" - # Note that redundant-returns-doc is broken for pylint 1.7.x for abstract - # methods. - # pylint: disable=redundant-returns-doc - def __init__(self, ascii_codepage='cp1252', key_path_prefix=''): """Initializes a Windows Registry file. @@ -86,10 +82,6 @@ def SetKeyPathPrefix(self, key_path_prefix): class WinRegistryFileReader(object): """Windows Registry file reader interface.""" - # Note that redundant-returns-doc is broken for pylint 1.7.x for abstract - # methods. - # pylint: disable=redundant-returns-doc - @abc.abstractmethod def Open(self, path, ascii_codepage='cp1252'): """Opens a Windows Registry file specified by the path. @@ -109,10 +101,6 @@ def Open(self, path, ascii_codepage='cp1252'): class WinRegistryKey(object): """Windows Registry key interface.""" - # Note that redundant-returns-doc and redundant-yields-doc are broken for - # pylint 1.7.x for abstract methods. - # pylint: disable=redundant-returns-doc,redundant-yields-doc - def __init__(self, key_path=''): """Initializes a Windows Registry key. @@ -235,10 +223,6 @@ def RecurseKeys(self): class WinRegistryValue(object): """Windows Registry value interface.""" - # Note that redundant-returns-doc is broken for pylint 1.7.x for abstract - # methods. - # pylint: disable=redundant-returns-doc - _DATA_TYPE_STRINGS = { 0: 'REG_NONE', 1: 'REG_SZ', diff --git a/dfwinreg/regf.py b/dfwinreg/regf.py index 1c9851e..ca8faa6 100644 --- a/dfwinreg/regf.py +++ b/dfwinreg/regf.py @@ -28,12 +28,37 @@ def __init__( """ super(REGFWinRegistryFile, self).__init__( ascii_codepage=ascii_codepage, key_path_prefix=key_path_prefix) - self._current_control_set_key_path = None self._emulate_virtual_keys = emulate_virtual_keys self._file_object = None self._regf_file = pyregf.file() self._regf_file.set_ascii_codepage(ascii_codepage) self._virtual_keys_by_path = {} + self._virtual_subkeys_by_parent = {} + + def _CreateKey(self, key_path, relative_key_path, pyregf_key): + """Creates a Windows Registry key. + + Args: + key_path (str): Windows Registry key path. + relative_key_path (str): Windows Registry key path relative to the file. + pyregf_key (pyregf.key): pyregf key object. + + Returns: + WinRegistryKey: patched Windows Registry key. + """ + virtual_subkeys = self._virtual_subkeys_by_parent.get( + relative_key_path, None) + if not virtual_subkeys: + return REGFWinRegistryKey(pyregf_key, key_path=key_path) + + name = relative_key_path.rsplit('\\', maxsplit=1)[-1] + registry_key = VirtualREGFWinRegistryKey( + name, pyregf_key, key_path=key_path) + + for name, pyregf_subkey in virtual_subkeys or []: + registry_key.AddVirtualSubKey(name, pyregf_subkey) + + return registry_key def _GetCurrentControlSetKeyPath(self): """Retrieves the key path of the current control set key. @@ -66,58 +91,41 @@ def _GetCurrentControlSetKeyPath(self): return f'\\ControlSet{control_set:03d}' - def _GetKeyByPathFromFile(self, key_path): + def _GetKeyByPathFromFile(self, relative_key_path): """Retrieves the key for a specific path form the Windows Registry file. Args: - key_path (str): Windows Registry key path relative to the file. + relative_key_path (str): Windows Registry key path relative to the file. Returns: pyregf.key: Registry key or None if not available. """ try: - return self._regf_file.get_key_by_path(key_path) + return self._regf_file.get_key_by_path(relative_key_path) except IOError: return None - def _GetVirtualKeyByPath(self, key_path): - """Retrieves a virtual key for a specific path. - - Args: - key_path (str): Windows Registry key path. - - Raises: - VirtualREGFWinRegistryKey: virtual key or None if not available. - """ - lookup_key_path = key_path.upper() - return self._virtual_keys_by_path.get(lookup_key_path, None) - - def AddVirtualKey(self, key_path, key): + def AddVirtualKey(self, relative_key_path, pyregf_key): """Adds a virtual key. Args: - key_path (str): key path of the virtual Windows Registry key. - key (pyregf.key): pyregf key object of the key. + relative_key_path (str): Windows Registry key path relative to the file. + pyregf_key (pyregf.key): pyregf key object of the key. Raises: - ValueError: if the key already exists. + ValueError: if the virtual key already exists. """ - key_path_upper = key_path.upper() - if key_path_upper.startswith(self._key_path_prefix_upper): - relative_key_path = key_path[self._key_path_prefix_length:] - elif key_path.startswith(definitions.KEY_PATH_SEPARATOR): - relative_key_path = key_path - key_path = ''.join([self._key_path_prefix, key_path]) - else: - raise ValueError(f'Unsupported key path: {key_path:s}') - lookup_key_path = relative_key_path.upper() if lookup_key_path in self._virtual_keys_by_path: - raise ValueError(f'Key: {key_path:s} already set') + raise ValueError(f'Key: {relative_key_path:s} already set') + + self._virtual_keys_by_path[lookup_key_path] = ( + relative_key_path, pyregf_key) - name = key_path.rsplit('\\', maxsplit=1)[-1] - self._virtual_keys_by_path[lookup_key_path] = VirtualREGFWinRegistryKey( - name, key, key_path=key_path) + parent_key_path, name = relative_key_path.rsplit('\\', maxsplit=1) + if parent_key_path not in self._virtual_subkeys_by_parent: + self._virtual_subkeys_by_parent[parent_key_path] = [] + self._virtual_subkeys_by_parent[parent_key_path].append((name, pyregf_key)) def Close(self): """Closes the Windows Registry file.""" @@ -145,36 +153,39 @@ def GetKeyByPath(self, key_path): if relative_key_path and relative_key_path[0] == '\\': relative_key_path = relative_key_path[1:] - relative_key_path_segments = relative_key_path.split('\\') - if not relative_key_path: - registry_key = self.GetRootKey() - - elif (self._emulate_virtual_keys and - relative_key_path_segments[0].upper() == 'CURRENTCONTROLSET'): - relative_key_path_segments.pop(0) - - virtual_key_path = '\\CurrentControlSet' - registry_key = self._GetVirtualKeyByPath(virtual_key_path) - if not registry_key: - current_control_set_key = self._GetKeyByPathFromFile( - self._current_control_set_key_path) - if current_control_set_key: - self.AddVirtualKey(virtual_key_path, current_control_set_key) - registry_key = self._GetVirtualKeyByPath(virtual_key_path) + return self.GetRootKey() - if relative_key_path_segments: - relative_sub_key_path = '\\'.join(relative_key_path_segments) - registry_key = registry_key.GetSubkeyByPath(relative_sub_key_path) - - else: - regf_key = self._GetKeyByPathFromFile(relative_key_path) - if not regf_key: - return None - - registry_key = REGFWinRegistryKey(regf_key, key_path=key_path) - - return registry_key + if self._emulate_virtual_keys: + relative_key_path_upper = relative_key_path.upper() + + lookup_key_path = None + relative_sub_key_path = None + for virtual_key_path in self._virtual_keys_by_path: + # Note that the virtual key path starts with a key path segment + # separator # but relative key path does not. + if relative_key_path_upper.startswith(virtual_key_path[1:]): + lookup_key_path = virtual_key_path + relative_sub_key_path = relative_key_path[len(virtual_key_path[1:]):] + break + + if lookup_key_path: + virtual_key_path, pyregf_key = self._virtual_keys_by_path.get( + lookup_key_path, None) + _, name = virtual_key_path.rsplit('\\', maxsplit=1) + virtual_key_path = ''.join([self._key_path_prefix, virtual_key_path]) + registry_key = VirtualREGFWinRegistryKey( + name, pyregf_key, key_path=virtual_key_path) + if not relative_sub_key_path: + return registry_key + + return registry_key.GetSubkeyByPath(relative_sub_key_path) + + pyregf_key = self._GetKeyByPathFromFile(relative_key_path) + if pyregf_key: + return self._CreateKey(key_path, relative_key_path, pyregf_key) + + return None def GetRootKey(self): """Retrieves the root key. @@ -182,24 +193,11 @@ def GetRootKey(self): Returns: WinRegistryKey: Windows Registry root key or None if not available. """ - regf_key = self._regf_file.get_root_key() - if not regf_key: + pyregf_key = self._regf_file.get_root_key() + if not pyregf_key: return None - if self._current_control_set_key_path: - current_control_set_key = self._GetKeyByPathFromFile( - self._current_control_set_key_path) - - registry_key = VirtualREGFWinRegistryKey( - '', regf_key, key_path=self._key_path_prefix) - registry_key.AddVirtualSubKey( - 'CurrentControlSet', current_control_set_key) - - else: - registry_key = REGFWinRegistryKey( - regf_key, key_path=self._key_path_prefix) - - return registry_key + return self._CreateKey(self._key_path_prefix, '', pyregf_key) def Open(self, file_object): """Opens the Windows Registry file using a file-like object. @@ -214,7 +212,11 @@ def Open(self, file_object): self._regf_file.open_file_object(self._file_object) if self._emulate_virtual_keys: - self._current_control_set_key_path = self._GetCurrentControlSetKeyPath() + key_path = self._GetCurrentControlSetKeyPath() + if key_path: + pyregf_key = self._GetKeyByPathFromFile(key_path) + if pyregf_key: + self.AddVirtualKey('\\CurrentControlSet', pyregf_key) return True @@ -410,7 +412,7 @@ def AddVirtualSubKey(self, name, subkey): subkey (pyregf.key): pyregf key object of the subkey. Raises: - ValueError: if the subkey already exists. + ValueError: if the virtual subkey already exists. """ lookup_name = name.upper() if lookup_name in self._virtual_subkeys_by_name: diff --git a/dfwinreg/registry.py b/dfwinreg/registry.py index 3c68c32..f27f0e9 100644 --- a/dfwinreg/registry.py +++ b/dfwinreg/registry.py @@ -284,7 +284,7 @@ def _GetFileByPath(self, key_path_upper): self.MapFile(key_path_prefix_upper, registry_file) # TODO: if HKEY_CURRENT_USER set 'HKEY_CURRENT_USER\\SOFTWARE\\CLASSES' as - #virtual key in the file. + # virtual key in the file. return key_path_prefix_upper, registry_file diff --git a/tests/regf.py b/tests/regf.py index b0ba30f..9e909a4 100644 --- a/tests/regf.py +++ b/tests/regf.py @@ -69,9 +69,9 @@ class REGFWinRegistryFileTest(test_lib.BaseTestCase): # pylint: disable=protected-access + # TODO: add tests for _CreateKey # TODO: add tests for _GetCurrentControlSetKeyPath # TODO: add tests for _GetKeyByPathFromFile - # TODO: add tests for _GetVirtualKeyByPath # TODO: add tests for AddVirtualKey def testOpenClose(self):