diff --git a/dfwinreg/creg.py b/dfwinreg/creg.py index 2bbfba7..d51e70b 100644 --- a/dfwinreg/creg.py +++ b/dfwinreg/creg.py @@ -54,10 +54,13 @@ def GetKeyByPath(self, key_path): creg_key = self._creg_file.get_key_by_path(relative_key_path) except IOError: creg_key = None + if not creg_key: return None - return CREGWinRegistryKey(creg_key, key_path=key_path) + return CREGWinRegistryKey( + creg_key, key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) def GetRootKey(self): """Retrieves the root key. @@ -69,7 +72,9 @@ def GetRootKey(self): if not creg_key: return None - return CREGWinRegistryKey(creg_key, key_path=self._key_path_prefix) + return CREGWinRegistryKey( + creg_key, key_path_prefix=self._key_path_prefix, + relative_key_path='') def Open(self, file_object): """Opens the Windows Registry file using a file-like object. @@ -88,14 +93,20 @@ def Open(self, file_object): class CREGWinRegistryKey(interface.WinRegistryKey): """Implementation of a Windows Registry key using pycreg.""" - def __init__(self, pycreg_key, key_path=''): - """Initializes a Windows Registry key object. + def __init__( + self, pycreg_key, key_helper=None, key_path_prefix='', + relative_key_path=''): + """Initializes a Windows Registry key. Args: pycreg_key (pycreg.key): pycreg key object. - key_path (Optional[str]): Windows Registry key path. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. + relative_key_path (Optional[str]): relative Windows Registry key path. """ - super(CREGWinRegistryKey, self).__init__(key_path=key_path) + super(CREGWinRegistryKey, self).__init__( + key_helper=key_helper, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) self._pycreg_key = pycreg_key @property @@ -144,8 +155,12 @@ def GetSubkeyByIndex(self, index): raise IndexError('Index out of bounds.') pycreg_key = self._pycreg_key.get_sub_key(index) - key_path = key_paths.JoinKeyPath([self._key_path, pycreg_key.name]) - return CREGWinRegistryKey(pycreg_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pycreg_key.name]) + return CREGWinRegistryKey( + pycreg_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) def GetSubkeyByName(self, name): """Retrieves a subkey by name. @@ -160,8 +175,12 @@ def GetSubkeyByName(self, name): if not pycreg_key: return None - key_path = key_paths.JoinKeyPath([self._key_path, pycreg_key.name]) - return CREGWinRegistryKey(pycreg_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pycreg_key.name]) + return CREGWinRegistryKey( + pycreg_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) def GetSubkeyByPath(self, key_path): """Retrieves a subkey by path. @@ -176,8 +195,12 @@ def GetSubkeyByPath(self, key_path): if not pycreg_key: return None - key_path = key_paths.JoinKeyPath([self._key_path, key_path]) - return CREGWinRegistryKey(pycreg_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, key_path]) + return CREGWinRegistryKey( + pycreg_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) def GetSubkeys(self): """Retrieves all subkeys within the key. @@ -186,8 +209,12 @@ def GetSubkeys(self): WinRegistryKey: Windows Registry subkey. """ for pycreg_key in self._pycreg_key.sub_keys: - key_path = key_paths.JoinKeyPath([self._key_path, pycreg_key.name]) - yield CREGWinRegistryKey(pycreg_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pycreg_key.name]) + yield CREGWinRegistryKey( + pycreg_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) def GetValueByName(self, name): """Retrieves a value by name. diff --git a/dfwinreg/fake.py b/dfwinreg/fake.py index e2cf251..19ff823 100644 --- a/dfwinreg/fake.py +++ b/dfwinreg/fake.py @@ -118,24 +118,29 @@ class FakeWinRegistryKey(interface.WinRegistryKey): """Fake implementation of a Windows Registry key.""" def __init__( - self, name, class_name=None, key_path='', last_written_time=None, - offset=None, subkeys=None, values=None): + self, name, class_name=None, key_helper=None, key_path_prefix='', + last_written_time=None, offset=None, relative_key_path='', subkeys=None, + values=None): """Initializes a Windows Registry key. Subkeys and values with duplicate names are silently ignored. Args: name (str): name of the Windows Registry key. - key_path (Optional[str]): Windows Registry key path. class_name (Optional[str]): class name of the Windows Registry key. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. last_written_time (Optional[int]): last written time, formatted as a FILETIME timestamp. offset (Optional[int]): offset of the key within the Windows Registry file. + relative_key_path (Optional[str]): relative Windows Registry key path. subkeys (Optional[list[FakeWinRegistryKey]]): list of subkeys. values (Optional[list[FakeWinRegistryValue]]): list of values. """ - super(FakeWinRegistryKey, self).__init__(key_path=key_path) + super(FakeWinRegistryKey, self).__init__( + key_helper=key_helper, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) self._class_name = class_name self._last_written_time = last_written_time self._name = name @@ -192,9 +197,12 @@ def _BuildKeyHierarchy(self, subkeys, values): continue self._subkeys[name] = registry_key + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, registry_key.name]) + # pylint: disable=protected-access - registry_key._key_path = key_paths.JoinKeyPath([ - self._key_path, registry_key.name]) + registry_key._key_path_prefix = self._key_path_prefix + registry_key._relative_key_path = relative_key_path if values: for registry_value in values: @@ -219,8 +227,11 @@ def AddSubkey(self, name, registry_key): self._subkeys[name_upper] = registry_key - key_path = key_paths.JoinKeyPath([self._key_path, name]) - registry_key._key_path = key_path # pylint: disable=protected-access + relative_key_path = key_paths.JoinKeyPath([self._relative_key_path, name]) + + # pylint: disable=protected-access + registry_key._key_path_prefix = self._key_path_prefix + registry_key._relative_key_path = relative_key_path def AddValue(self, registry_value): """Adds a value. diff --git a/dfwinreg/interface.py b/dfwinreg/interface.py index 1ab04d0..f12d513 100644 --- a/dfwinreg/interface.py +++ b/dfwinreg/interface.py @@ -101,14 +101,18 @@ def Open(self, path, ascii_codepage='cp1252'): class WinRegistryKey(object): """Windows Registry key interface.""" - def __init__(self, key_path=''): + def __init__(self, key_helper=None, key_path_prefix='', relative_key_path=''): """Initializes a Windows Registry key. Args: - key_path (Optional[str]): Windows Registry key path. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. + relative_key_path (Optional[str]): relative Windows Registry key path. """ super(WinRegistryKey, self).__init__() - self._key_path = key_paths.JoinKeyPath([key_path]) + self._key_helper = key_helper + self._key_path_prefix = key_path_prefix + self._relative_key_path = relative_key_path @property @abc.abstractmethod @@ -143,7 +147,8 @@ def offset(self): @property def path(self): """str: Windows Registry key path.""" - return self._key_path + return key_paths.JoinKeyPath([ + self._key_path_prefix, self._relative_key_path]) @abc.abstractmethod def GetSubkeyByIndex(self, index): @@ -220,6 +225,10 @@ def RecurseKeys(self): yield key +class WinRegistryKeyHelper(object): + """Windows Registry key helper.""" + + class WinRegistryValue(object): """Windows Registry value interface.""" diff --git a/dfwinreg/regf.py b/dfwinreg/regf.py index ca8faa6..d32c9be 100644 --- a/dfwinreg/regf.py +++ b/dfwinreg/regf.py @@ -29,36 +29,10 @@ def __init__( super(REGFWinRegistryFile, self).__init__( ascii_codepage=ascii_codepage, key_path_prefix=key_path_prefix) self._emulate_virtual_keys = emulate_virtual_keys + self._key_helper = REGFWinRegistryKeyHelper() 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. @@ -115,17 +89,7 @@ def AddVirtualKey(self, relative_key_path, pyregf_key): Raises: ValueError: if the virtual key already exists. """ - lookup_key_path = relative_key_path.upper() - if lookup_key_path in self._virtual_keys_by_path: - raise ValueError(f'Key: {relative_key_path:s} already set') - - self._virtual_keys_by_path[lookup_key_path] = ( - relative_key_path, pyregf_key) - - 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)) + self._key_helper.AddVirtualKey(relative_key_path, pyregf_key) def Close(self): """Closes the Windows Registry file.""" @@ -146,7 +110,7 @@ def GetKeyByPath(self, key_path): 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]) + key_path = ''.join([self._key_path_prefix, relative_key_path]) else: return None @@ -157,35 +121,14 @@ def GetKeyByPath(self, key_path): return self.GetRootKey() 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) + registry_key = self._key_helper.GetKeyByPath( + self._key_path_prefix, relative_key_path) + if registry_key: + return registry_key pyregf_key = self._GetKeyByPathFromFile(relative_key_path) - if pyregf_key: - return self._CreateKey(key_path, relative_key_path, pyregf_key) - - return None + return self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) def GetRootKey(self): """Retrieves the root key. @@ -194,10 +137,7 @@ def GetRootKey(self): WinRegistryKey: Windows Registry root key or None if not available. """ pyregf_key = self._regf_file.get_root_key() - if not pyregf_key: - return None - - return self._CreateKey(self._key_path_prefix, '', pyregf_key) + return self._key_helper.CreateKey(self._key_path_prefix, '', pyregf_key) def Open(self, file_object): """Opens the Windows Registry file using a file-like object. @@ -224,14 +164,20 @@ def Open(self, file_object): class REGFWinRegistryKey(interface.WinRegistryKey): """Implementation of a Windows Registry key using pyregf.""" - def __init__(self, pyregf_key, key_path=''): + def __init__( + self, pyregf_key, key_helper=None, key_path_prefix='', + relative_key_path=''): """Initializes a Windows Registry key. Args: pyregf_key (pyregf.key): pyregf key object. - key_path (Optional[str]): Windows Registry key path. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. + relative_key_path (Optional[str]): relative Windows Registry key path. """ - super(REGFWinRegistryKey, self).__init__(key_path=key_path) + super(REGFWinRegistryKey, self).__init__( + key_helper=key_helper, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) self._pyregf_key = pyregf_key @property @@ -284,8 +230,10 @@ def GetSubkeyByIndex(self, index): raise IndexError('Index out of bounds.') pyregf_key = self._pyregf_key.get_sub_key(index) - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - return REGFWinRegistryKey(pyregf_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pyregf_key.name]) + return self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) def GetSubkeyByName(self, name): """Retrieves a subkey by name. @@ -300,8 +248,10 @@ def GetSubkeyByName(self, name): if not pyregf_key: return None - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - return REGFWinRegistryKey(pyregf_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pyregf_key.name]) + return self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) def GetSubkeyByPath(self, key_path): """Retrieves a subkey by path. @@ -316,8 +266,10 @@ def GetSubkeyByPath(self, key_path): if not pyregf_key: return None - key_path = key_paths.JoinKeyPath([self._key_path, key_path]) - return REGFWinRegistryKey(pyregf_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, key_path]) + return self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) def GetSubkeys(self): """Retrieves all subkeys within the key. @@ -326,8 +278,10 @@ def GetSubkeys(self): WinRegistryKey: Windows Registry subkey. """ for pyregf_key in self._pyregf_key.sub_keys: - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - yield REGFWinRegistryKey(pyregf_key, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, pyregf_key.name]) + yield self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) def GetValues(self): """Retrieves all values within the key. @@ -361,16 +315,21 @@ class VirtualREGFWinRegistryKey(REGFWinRegistryKey): but do exist at run-time, like HKEY_LOCAL_MACHINE\\System\\CurrentControlSet. """ - def __init__(self, name, pyregf_key, key_path=''): + def __init__( + self, name, pyregf_key, key_helper=None, key_path_prefix='', + relative_key_path=''): """Initializes a virtual Windows Registry key. Args: name (str): name of the Windows Registry key. pyregf_key (pyregf.key): pyregf key object. - key_path (Optional[str]): Windows Registry key path. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. + relative_key_path (Optional[str]): relative Windows Registry key path. """ super(VirtualREGFWinRegistryKey, self).__init__( - pyregf_key, key_path=key_path) + pyregf_key, key_helper=key_helper, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) self._name = name self._virtual_subkeys = [] self._virtual_subkeys_by_name = {} @@ -433,23 +392,21 @@ def GetSubkeyByIndex(self, index): Raises: IndexError: if the index is out of bounds. """ - if index < 0 or index >= self.number_of_subkeys: + if index >= self.number_of_subkeys: raise IndexError('Index out of bounds.') - if index < self._pyregf_key.number_of_sub_keys: - pyregf_key = self._pyregf_key.get_sub_key(index) - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - subkey = REGFWinRegistryKey(pyregf_key, key_path=key_path) - - else: + if index >= self._pyregf_key.number_of_sub_keys: index -= self._pyregf_key.number_of_sub_keys virtual_name, virtual_subkey = self._virtual_subkeys[index] - key_path = key_paths.JoinKeyPath([self._key_path, virtual_name]) - subkey = VirtualREGFWinRegistryKey( - virtual_name, virtual_subkey, key_path=key_path) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, virtual_name]) + return VirtualREGFWinRegistryKey( + virtual_name, virtual_subkey, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) - return subkey + return super(VirtualREGFWinRegistryKey, self).GetSubkeyByIndex(index) def GetSubkeyByName(self, name): """Retrieves a subkey by name. @@ -462,16 +419,14 @@ def GetSubkeyByName(self, name): """ virtual_name, virtual_sub_key = self._GetVirtualSubKeyByName(name) if virtual_sub_key: - key_path = key_paths.JoinKeyPath([self._key_path, virtual_name]) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, virtual_name]) return VirtualREGFWinRegistryKey( - virtual_name, virtual_sub_key, key_path=key_path) + virtual_name, virtual_sub_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) - pyregf_key = self._pyregf_key.get_sub_key_by_name(name) - if not pyregf_key: - return None - - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - return REGFWinRegistryKey(pyregf_key, key_path=key_path) + return super(VirtualREGFWinRegistryKey, self).GetSubkeyByName(name) def GetSubkeyByPath(self, key_path): """Retrieves a subkey by path. @@ -489,24 +444,26 @@ def GetSubkeyByPath(self, key_path): virtual_name, virtual_sub_key = self._GetVirtualSubKeyByName( key_path_segments[0]) - if not virtual_sub_key: - pyregf_key = self._pyregf_key.get_sub_key_by_path(key_path) - else: + if virtual_sub_key: key_path_segments.pop(0) if not key_path_segments: - key_path = key_paths.JoinKeyPath([self._key_path, virtual_name]) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, virtual_name]) return VirtualREGFWinRegistryKey( - virtual_name, virtual_sub_key, key_path=key_path) + virtual_name, virtual_sub_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) sub_key_path = '\\'.join(key_path_segments) pyregf_key = virtual_sub_key.get_sub_key_by_path(sub_key_path) - if not pyregf_key: - return None + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, key_path]) + return self._key_helper.CreateKey( + self._key_path_prefix, relative_key_path, pyregf_key) - key_path = key_paths.JoinKeyPath([self._key_path, key_path]) - return REGFWinRegistryKey(pyregf_key, key_path=key_path) + return super(VirtualREGFWinRegistryKey, self).GetSubkeyByPath(key_path) def GetSubkeys(self): """Retrieves all subkeys within the key. @@ -514,14 +471,122 @@ def GetSubkeys(self): Yields: WinRegistryKey: Windows Registry subkey. """ - for pyregf_key in self._pyregf_key.sub_keys: - key_path = key_paths.JoinKeyPath([self._key_path, pyregf_key.name]) - yield REGFWinRegistryKey(pyregf_key, key_path=key_path) + yield from super(VirtualREGFWinRegistryKey, self).GetSubkeys() for virtual_name, virtual_sub_key in self._virtual_subkeys: - key_path = key_paths.JoinKeyPath([self._key_path, virtual_name]) + relative_key_path = key_paths.JoinKeyPath([ + self._relative_key_path, virtual_name]) yield VirtualREGFWinRegistryKey( - virtual_name, virtual_sub_key, key_path=key_path) + virtual_name, virtual_sub_key, key_helper=self._key_helper, + key_path_prefix=self._key_path_prefix, + relative_key_path=relative_key_path) + + +class REGFWinRegistryKeyHelper(interface.WinRegistryKeyHelper): + """Windows Registry key helper.""" + + def __init__(self): + """Initializes the Windows Registry key helper.""" + super(REGFWinRegistryKeyHelper, self).__init__() + self._virtual_keys_by_path = {} + self._virtual_subkeys_by_parent = {} + + def AddVirtualKey(self, relative_key_path, pyregf_key): + """Adds a virtual key. + + Args: + relative_key_path (str): Windows Registry key path relative to the file, + with a leading key path segment separator. + pyregf_key (pyregf.key): pyregf key object of the key. + + Raises: + ValueError: if the virtual key already exists. + """ + lookup_key_path = relative_key_path.upper() + if lookup_key_path in self._virtual_keys_by_path: + raise ValueError(f'Key: {relative_key_path:s} already set') + + self._virtual_keys_by_path[lookup_key_path] = ( + relative_key_path, pyregf_key) + + parent_key_path, name = relative_key_path.rsplit('\\', maxsplit=1) + + lookup_key_path = parent_key_path.upper() + if lookup_key_path not in self._virtual_subkeys_by_parent: + self._virtual_subkeys_by_parent[lookup_key_path] = [] + self._virtual_subkeys_by_parent[lookup_key_path].append((name, pyregf_key)) + + def CreateKey(self, key_path_prefix, relative_key_path, pyregf_key): + """Creates a Windows Registry key. + + Args: + key_path_prefix (str): Windows Registry key path prefix. + relative_key_path (str): Windows Registry key path relative to the file, + with a leading key path segment separator. + pyregf_key (pyregf.key): pyregf key object. + + Returns: + WinRegistryKey: Windows Registry key or None if pyregf key object is not + set. + """ + if not pyregf_key: + return None + + lookup_key_path = relative_key_path.upper() + virtual_subkeys = self._virtual_subkeys_by_parent.get(lookup_key_path, None) + + if not virtual_subkeys: + return REGFWinRegistryKey( + pyregf_key, key_helper=self, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) + + name = relative_key_path.rsplit('\\', maxsplit=1)[-1] + registry_key = VirtualREGFWinRegistryKey( + name, pyregf_key, key_helper=self, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) + + for name, pyregf_subkey in virtual_subkeys or []: + registry_key.AddVirtualSubKey(name, pyregf_subkey) + + return registry_key + + def GetKeyByPath(self, key_path_prefix, relative_key_path): + """Retrieves a key. + + Args: + key_path_prefix (str): Windows Registry key path prefix. + relative_key_path (str): Windows Registry key path relative to the file, + without a leading key path segment separator. + + Returns: + WinRegistryKey: Windows Registry key or None if not found. + """ + lookup_key_path = None + relative_sub_key_path = None + + # TODO: use scan tree of path segments for faster lookup. + relative_key_path_upper = relative_key_path.upper() + 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 not lookup_key_path: + return None + + virtual_key_path, pyregf_key = self._virtual_keys_by_path.get( + lookup_key_path, None) + _, name = virtual_key_path.rsplit('\\', maxsplit=1) + registry_key = VirtualREGFWinRegistryKey( + name, pyregf_key, key_helper=self, key_path_prefix=key_path_prefix, + relative_key_path=virtual_key_path[1:]) + if not relative_sub_key_path: + return registry_key + + return registry_key.GetSubkeyByPath(relative_sub_key_path) class REGFWinRegistryValue(interface.WinRegistryValue): diff --git a/dfwinreg/registry.py b/dfwinreg/registry.py index f27f0e9..fe2c44c 100644 --- a/dfwinreg/registry.py +++ b/dfwinreg/registry.py @@ -3,6 +3,7 @@ from dfwinreg import definitions from dfwinreg import key_paths +from dfwinreg import regf from dfwinreg import virtual @@ -111,6 +112,8 @@ class WinRegistry(object): 'HKEY_LOCAL_MACHINE\\Software\\Microsoft\\Windows NT\\CurrentVersion\\' 'ProfileList') + _USER_SOFTWARE_CLASSES_KEY_PATH = 'HKEY_CURRENT_USER\\Software\\Classes' + # TODO: add support for HKEY_CLASSES_ROOT # TODO: add support for HKEY_CURRENT_CONFIG # TODO: add support for HKEY_CURRENT_USER @@ -286,6 +289,14 @@ def _GetFileByPath(self, key_path_upper): # TODO: if HKEY_CURRENT_USER set 'HKEY_CURRENT_USER\\SOFTWARE\\CLASSES' as # virtual key in the file. + if key_path_prefix_upper == 'HKEY_CURRENT_USER' and isinstance( + registry_file, regf.REGFWinRegistryFile): + registry_key = self.GetKeyByPath(self._USER_SOFTWARE_CLASSES_KEY_PATH) + if registry_key: + # pylint: disable=protected-access + pyregf_key = registry_key._pyregf_key + registry_file.AddVirtualKey('\\Software\\Classes', pyregf_key) + return key_path_prefix_upper, registry_file def _GetKeyByPathFromFile(self, key_path): @@ -332,7 +343,7 @@ def _GetRootVirtualKey(self): sub_registry_key = self.GetKeyByPath(sub_key_name) if not sub_registry_key: sub_registry_key = virtual.VirtualWinRegistryKey( - sub_key_name, key_path=sub_key_name, registry=self) + sub_key_name, registry=self, relative_key_path=sub_key_name) if sub_key_name == 'HKEY_LOCAL_MACHINE': local_machine_key = sub_registry_key @@ -346,7 +357,7 @@ def _GetRootVirtualKey(self): sub_registry_key = self.GetKeyByPath(sub_key_path) if not sub_registry_key: sub_registry_key = virtual.VirtualWinRegistryKey( - sub_key_name, key_path=sub_key_path, registry=self) + sub_key_name, registry=self, relative_key_path=sub_key_name) local_machine_key.AddSubkey(sub_key_name, sub_registry_key) @@ -505,7 +516,6 @@ def GetKeyByPath(self, key_path): registry_key = self._GetKeyByPathFromFile(key_path) if not registry_key: - # TODO: handle NTUSER.DAT not having SOFTWARE\\CLASSES key. registry_key = self._GetVirtualKeyByPath(key_path) return registry_key diff --git a/dfwinreg/registry_searcher.py b/dfwinreg/registry_searcher.py index 9044236..b53db2e 100644 --- a/dfwinreg/registry_searcher.py +++ b/dfwinreg/registry_searcher.py @@ -152,7 +152,7 @@ def _CheckKeyPath(self, registry_key, search_depth): self._key_path_segments[search_depth - 1] = segment_name if search_depth > 0: - key_name = registry_key.path.split('\\')[-1] + key_name = registry_key.path.rsplit('\\', maxsplit=1)[-1] if self._is_regex: # pylint: disable=no-member if not segment_name.match(key_name): @@ -278,7 +278,7 @@ def CompareNameWithKeyPathSegment(self, registry_key, segment_index): that of the find specification, False if not or if the find specification has no key path defined. """ - key_name = registry_key.path.split('\\')[-1] + key_name = registry_key.path.rsplit('\\', maxsplit=1)[-1] return self._CompareWithKeyPathSegment(key_name, segment_index) def HasKeyPath(self): diff --git a/dfwinreg/virtual.py b/dfwinreg/virtual.py index 9935d1b..0bc774c 100644 --- a/dfwinreg/virtual.py +++ b/dfwinreg/virtual.py @@ -1,8 +1,6 @@ # -*- coding: utf-8 -*- """Virtual Windows Registry key implementation.""" -import collections - from dfwinreg import definitions from dfwinreg import interface from dfwinreg import key_paths @@ -17,19 +15,27 @@ class VirtualWinRegistryKey(interface.WinRegistryKey): HKEY_LOCAL_MACHINE\\System. """ - def __init__(self, name, key_path='', registry=None): + # TODO: move registry to key_helper + def __init__( + self, name, key_helper=None, key_path_prefix='', registry=None, + relative_key_path=''): """Initializes a Windows Registry key. Args: name (str): name of the Windows Registry key. - key_path (Optional[str]): Windows Registry key path. + key_helper (Optional[WinRegistryKeyHelper]): Windows Registry key helper. + key_path_prefix (Optional[str]): Windows Registry key path prefix. registry (Optional[WinRegistry]): Windows Registry. + relative_key_path (Optional[str]): relative Windows Registry key path. """ - super(VirtualWinRegistryKey, self).__init__(key_path=key_path) + super(VirtualWinRegistryKey, self).__init__( + key_helper=key_helper, key_path_prefix=key_path_prefix, + relative_key_path=relative_key_path) self._name = name self._registry = registry self._registry_key = None - self._subkeys = collections.OrderedDict() + self._subkeys = [] + self._subkeys_by_name = {} @property def class_name(self): @@ -93,8 +99,11 @@ def _GetKeyFromRegistry(self): if not self._registry: return + key_path = key_paths.JoinKeyPath([ + self._key_path_prefix, self._relative_key_path]) + try: - self._registry_key = self._registry.GetKeyByPath(self._key_path) + self._registry_key = self._registry.GetKeyByPath(key_path) except RuntimeError: pass @@ -143,13 +152,17 @@ def AddSubkey(self, name, registry_key): KeyError: if the subkey already exists. """ name_upper = name.upper() - if name_upper in self._subkeys: + if name_upper in self._subkeys_by_name: raise KeyError(f'Subkey: {name:s} already exists.') - self._subkeys[name_upper] = registry_key + self._subkeys_by_name[name_upper] = len(self._subkeys) + self._subkeys.append(registry_key) + + relative_key_path = self._JoinKeyPath([self._relative_key_path, name]) - key_path = self._JoinKeyPath([self._key_path, name]) - registry_key._key_path = key_path # pylint: disable=protected-access + # pylint: disable=protected-access + registry_key._key_path_prefix = self._key_path_prefix + registry_key._relative_key_path = relative_key_path def GetSubkeyByIndex(self, index): """Retrieves a subkey by index. @@ -166,12 +179,10 @@ def GetSubkeyByIndex(self, index): if not self._registry_key and self._registry: self._GetKeyFromRegistry() - subkeys = list(self._subkeys.values()) - - if index < 0 or index >= len(subkeys): + if index < 0 or index >= len(self._subkeys): raise IndexError('Index out of bounds.') - return subkeys[index] + return self._subkeys[index] def GetSubkeyByName(self, name): """Retrieves a subkey by name. @@ -185,7 +196,11 @@ def GetSubkeyByName(self, name): if not self._registry_key and self._registry: self._GetKeyFromRegistry() - return self._subkeys.get(name.upper(), None) + index = self._subkeys_by_name.get(name.upper(), None) + if index is None: + return None + + return self._subkeys[index] def GetSubkeyByPath(self, key_path): """Retrieves a subkey by path. @@ -210,13 +225,13 @@ def GetSubkeyByPath(self, key_path): def GetSubkeys(self): """Retrieves all subkeys within the key. - Returns: - generator[WinRegistryKey]: Windows Registry subkey generator. + Yields: + WinRegistryKey: Windows Registry subkey. """ if not self._registry_key and self._registry: self._GetKeyFromRegistry() - return iter(self._subkeys.values()) + yield from self._subkeys def GetValueByName(self, name): """Retrieves a value by name. @@ -238,13 +253,11 @@ def GetValueByName(self, name): def GetValues(self): """Retrieves all values within the key. - Returns: - generator[WinRegistryValue]: Windows Registry value generator. + Yields: + WinRegistryValue: Windows Registry value. """ if not self._registry_key and self._registry: self._GetKeyFromRegistry() if self._registry_key: - return self._registry_key.GetValues() - - return iter([]) + yield from self._registry_key.GetValues() diff --git a/tests/fake.py b/tests/fake.py index a31036c..0552971 100644 --- a/tests/fake.py +++ b/tests/fake.py @@ -143,6 +143,8 @@ class FakeWinRegistryKeyTest(test_lib.BaseTestCase): # pylint: disable=protected-access + _KEY_PATH_PREFIX = 'HKEY_CURRENT_USER\\Software' + def _CreateTestKey(self): """Creates a fake Windows Registry key for testing. @@ -150,8 +152,8 @@ def _CreateTestKey(self): FakeWinRegistryKey: fake Windows Registry key. """ registry_key = fake.FakeWinRegistryKey( - 'Software', key_path='HKEY_CURRENT_USER\\Software', - last_written_time=0) + 'Software', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='') sub_registry_key = fake.FakeWinRegistryKey( 'Microsoft', last_written_time=0) @@ -188,8 +190,8 @@ def testProperties(self): def testBuildKeyHierarchy(self): """Tests the BuildKeyHierarchy function.""" test_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft', - last_written_time=0) + 'Microsoft', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='Microsoft') test_value = fake.FakeWinRegistryValue('') @@ -198,15 +200,15 @@ def testBuildKeyHierarchy(self): # Test with subkeys and values. registry_key = fake.FakeWinRegistryKey( - 'Software', key_path='HKEY_CURRENT_USER\\Software', - last_written_time=0, + 'Software', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='', subkeys=[test_key], values=[test_value]) self.assertIsNotNone(registry_key) # Test with duplicate subkeys and values. registry_key = fake.FakeWinRegistryKey( - 'Software', key_path='HKEY_CURRENT_USER\\Software', - last_written_time=0, + 'Software', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='', subkeys=[test_key, test_key], values=[test_value, test_value]) self.assertIsNotNone(registry_key) @@ -214,12 +216,12 @@ def testBuildKeyHierarchy(self): def testAddSubkey(self): """Tests the AddSubkey function.""" registry_key = fake.FakeWinRegistryKey( - 'Software', key_path='HKEY_CURRENT_USER\\Software', - last_written_time=0) + 'Software', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='') sub_registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft', - last_written_time=0) + 'Microsoft', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='Microsoft') registry_key.AddSubkey(sub_registry_key.name, sub_registry_key) @@ -229,8 +231,8 @@ def testAddSubkey(self): def testAddValue(self): """Tests the AddValue function.""" registry_key = fake.FakeWinRegistryKey( - 'Software', key_path='HKEY_CURRENT_USER\\Software', - last_written_time=0) + 'Software', key_path_prefix=self._KEY_PATH_PREFIX, + last_written_time=0, relative_key_path='') registry_value = fake.FakeWinRegistryValue('') diff --git a/tests/registry_searcher.py b/tests/registry_searcher.py index d53d79f..c9796f9 100644 --- a/tests/registry_searcher.py +++ b/tests/registry_searcher.py @@ -66,7 +66,8 @@ def testCheckKeyPath(self): key_path='HKEY_CURRENT_USER\\Software\\Microsoft') registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') result = find_spec._CheckKeyPath(registry_key, 3) self.assertTrue(result) @@ -90,7 +91,8 @@ def testCheckKeyPath(self): key_path_regex=['HKEY_CURRENT_USER', 'Software', 'Microsoft']) registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') result = find_spec._CheckKeyPath(registry_key, 3) self.assertTrue(result) @@ -105,7 +107,8 @@ def testCheckKeyPath(self): key_path_regex=['HKEY_CURRENT_USER', 'Software', 'Mi(rosoft']) registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') result = find_spec._CheckKeyPath(registry_key, 3) self.assertFalse(result) @@ -160,7 +163,8 @@ def testAtMaximumDepth(self): def testCompareKeyPath(self): """Test the CompareKeyPath function.""" registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') find_spec = registry_searcher.FindSpec( key_path='HKEY_CURRENT_USER\\Software\\Microsoft') @@ -177,7 +181,8 @@ def testCompareKeyPath(self): def testCompareNameWithKeyPathSegment(self): """Test the CompareNameWithKeyPathSegment function.""" registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') find_spec = registry_searcher.FindSpec( key_path='HKEY_CURRENT_USER\\Software\\Microsoft') @@ -227,7 +232,8 @@ def testMatches(self): key_path='HKEY_CURRENT_USER\\Software\\Microsoft') registry_key = fake.FakeWinRegistryKey( - 'Microsoft', key_path='HKEY_CURRENT_USER\\Software\\Microsoft') + 'Microsoft', key_path_prefix='HKEY_CURRENT_USER\\Software', + relative_key_path='Microsoft') result = find_spec.Matches(registry_key, 3) self.assertEqual(result, (True, True)) @@ -294,8 +300,7 @@ def testFind(self): # Test with key path regular expression. find_spec = registry_searcher.FindSpec( - key_path_regex=[ - 'HKEY_LOCAL_MACHINE', 'System', 'ControlSet001', '.*']) + key_path_regex=['HKEY_LOCAL_MACHINE', 'System', 'ControlSet001', '.*']) expected_key_paths = [ 'HKEY_LOCAL_MACHINE\\System\\ControlSet001\\Control', diff --git a/tests/virtual.py b/tests/virtual.py index a1b640e..aa3933e 100644 --- a/tests/virtual.py +++ b/tests/virtual.py @@ -63,7 +63,7 @@ def _CreateTestKey(self): VirtualWinRegistryKey: virtual Windows Registry key. """ registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='HKEY_LOCAL_MACHINE') + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE') sub_registry_key = virtual.VirtualWinRegistryKey('System') registry_key.AddSubkey(sub_registry_key.name, sub_registry_key) @@ -149,23 +149,25 @@ def testPropertiesWithMappedRegistry(self): def testGetKeyFromRegistry(self): """Tests the _GetKeyFromRegistry function.""" registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='') + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE') registry_key._GetKeyFromRegistry() test_win_registry = TestWinRegistry() registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='', registry=test_win_registry) + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE', + registry=test_win_registry) registry_key._GetKeyFromRegistry() test_win_registry = ErrorWinRegistry() registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='', registry=test_win_registry) + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE', + registry=test_win_registry) registry_key._GetKeyFromRegistry() def testJoinKeyPath(self): """Tests the _JoinKeyPath function.""" registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='') + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE') expected_path = 'HKEY_LOCAL_MACHINE\\Software' path = registry_key._JoinKeyPath(['HKEY_LOCAL_MACHINE', 'Software']) @@ -174,10 +176,10 @@ def testJoinKeyPath(self): def testAddSubkey(self): """Tests the AddSubkey function.""" registry_key = virtual.VirtualWinRegistryKey( - 'HKEY_LOCAL_MACHINE', key_path='') + 'HKEY_LOCAL_MACHINE', relative_key_path='HKEY_LOCAL_MACHINE') sub_registry_key = virtual.VirtualWinRegistryKey( - 'System', key_path='HKEY_LOCAL_MACHINE') + 'System', relative_key_path='HKEY_LOCAL_MACHINE') registry_key.AddSubkey(sub_registry_key.name, sub_registry_key)