From d19932b6e054ff904ae55037b25ba825085da57c Mon Sep 17 00:00:00 2001 From: Joachim Metz Date: Mon, 26 Feb 2024 21:13:27 +0100 Subject: [PATCH] Changes for mapping HKEY_CURRENT_USER\Software\Classes (#222) --- dfwinreg/registry.py | 70 +++++++++++++++++---------- tests/registry.py | 111 ++++++++++++++++++++++++++++--------------- 2 files changed, 118 insertions(+), 63 deletions(-) diff --git a/dfwinreg/registry.py b/dfwinreg/registry.py index a6e2d7f..c2e6792 100644 --- a/dfwinreg/registry.py +++ b/dfwinreg/registry.py @@ -81,16 +81,14 @@ class WinRegistry(object): WinRegistryFileMapping( 'HKEY_LOCAL_MACHINE\\System', '%SystemRoot%\\System32\\config\\SYSTEM', - ['\\MountedDevices', '\\Select', '\\Setup']), - ] + ['\\MountedDevices', '\\Select', '\\Setup'])] _ROOT_KEY_ALIASES = { 'HKCC': 'HKEY_CURRENT_CONFIG', 'HKCR': 'HKEY_CLASSES_ROOT', 'HKCU': 'HKEY_CURRENT_USER', 'HKLM': 'HKEY_LOCAL_MACHINE', - 'HKU': 'HKEY_USERS', - } + 'HKU': 'HKEY_USERS'} _ROOT_SUB_KEYS = frozenset([ 'HKEY_CLASSES_ROOT', @@ -99,15 +97,13 @@ class WinRegistry(object): 'HKEY_DYN_DATA', 'HKEY_LOCAL_MACHINE', 'HKEY_PERFORMANCE_DATA', - 'HKEY_USERS', - ]) + 'HKEY_USERS']) _LOCAL_MACHINE_SUB_KEYS = frozenset([ 'SAM', 'Security', 'Software', - 'System', - ]) + 'System']) _SELECT_KEY_PATH = 'HKEY_LOCAL_MACHINE\\System\\Select' @@ -220,28 +216,38 @@ def _GetFileByPath(self, key_path_upper): a resolved root key alias. Returns: - tuple: consists: - - str: upper case key path prefix - WinRegistryFile: corresponding Windows Registry file or None if not - available. + tuple[st, WinRegistryFile]: upper case key path prefix and corresponding + Windows Registry file or None if not available. """ - key_path_prefix, registry_file = self._GetCachedFileByPath(key_path_upper) - if registry_file: - return key_path_prefix, registry_file + key_path_prefix, cached_registry_file = self._GetCachedFileByPath( + key_path_upper) + + # NTUSER.DAT and UsrClass.dat have overlapping keys hence for key paths that + # start "HKEY_CURRENT_USER\\Software\\Classes" do not assume the Windows + # Registry file for "HKEY_CURRENT_USER" contains the classes key as well. + + if cached_registry_file and ( + key_path_prefix != 'HKEY_CURRENT_USER' or + not key_path_upper.startswith('HKEY_CURRENT_USER\\SOFTWARE\\CLASSES')): + return key_path_prefix, cached_registry_file key_path_prefix = '' candidate_mappings = self._GetCandidateFileMappingsByPath(key_path_upper) matching_mappings = [] for mapping in candidate_mappings: + mapping_key_path_prefix = mapping.key_path_prefix.upper() + if mapping_key_path_prefix in self._registry_files: + # Skip the Windows Registry file if it is already cached. + continue + try: registry_file = self._OpenFile(mapping.windows_path) except IOError: registry_file = None - if len(key_path_prefix) < len(mapping.key_path_prefix): - key_path_prefix = mapping.key_path_prefix.upper() + if len(key_path_prefix) < len(mapping_key_path_prefix): + key_path_prefix = mapping_key_path_prefix if not registry_file: continue @@ -257,16 +263,27 @@ def _GetFileByPath(self, key_path_upper): if match: matching_mappings.append((mapping, registry_file)) - if not matching_mappings or len(matching_mappings) > 1: - for _, registry_file in matching_mappings: - registry_file.Close() - return key_path_prefix, None + if not matching_mappings: + return key_path_prefix, cached_registry_file - mapping, registry_file = matching_mappings[0] - key_path_prefix = mapping.key_path_prefix.upper() - self.MapFile(key_path_prefix, registry_file) + key_path_prefix = None + registry_file = None + + while matching_mappings: + mapping, matching_registry_file = matching_mappings.pop(0) + + if (not key_path_prefix or + len(mapping.key_path_prefix) > len(key_path_prefix)): + if registry_file: + registry_file.Close() + + key_path_prefix = mapping.key_path_prefix + registry_file = matching_registry_file + + key_path_prefix_upper = key_path_prefix.upper() + self.MapFile(key_path_prefix_upper, registry_file) - return key_path_prefix, registry_file + return key_path_prefix_upper, registry_file def _GetKeyByPathFromFile(self, key_path): """Retrieves the key for a specific path from file. @@ -485,6 +502,7 @@ 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/tests/registry.py b/tests/registry.py index 30dc01f..e4aacde 100644 --- a/tests/registry.py +++ b/tests/registry.py @@ -76,7 +76,7 @@ def _GetFileByPath(self, key_path_upper): class TestREGFWinRegistryFileReader(interface.WinRegistryFileReader): - """Single file Windows Registry file reader.""" + """Windows Registry file reader that reads a single test file.""" def __init__(self, emulate_virtual_keys=True): """Initializes a file Windows Registry file reader. @@ -87,6 +87,13 @@ def __init__(self, emulate_virtual_keys=True): """ super(TestREGFWinRegistryFileReader, self).__init__() self._emulate_virtual_keys = emulate_virtual_keys + self._file_objects = [] + + def __del__(self): + """Cleans up the Windows Registry file reader.""" + while self._file_objects: + file_object = self._file_objects.pop(0) + file_object.close() def Open(self, path, ascii_codepage='cp1252'): """Opens the Windows Registry file specified by the path. @@ -112,13 +119,30 @@ def Open(self, path, ascii_codepage='cp1252'): registry_file.Open(file_object) except IOError: file_object.close() + file_object = None registry_file = None + self._file_objects.append(file_object) + return registry_file class TestREGFWinRegistryFileReaderMapped(TestREGFWinRegistryFileReader): - """Single file Windows Registry file reader that maps Windows paths.""" + """Windows Registry file reader that maps Windows paths to test files.""" + + _MAPPED_PATHS = { + '%SystemRoot%\\System32\\config\\SAM': ( + os.path.join(test_lib.TEST_DATA_PATH, 'SAM')), + '%SystemRoot%\\System32\\config\\SECURITY': ( + os.path.join(test_lib.TEST_DATA_PATH, 'SECURITY')), + '%SystemRoot%\\System32\\config\\SOFTWARE': ( + os.path.join(test_lib.TEST_DATA_PATH, 'SOFTWARE')), + '%SystemRoot%\\System32\\config\\SYSTEM': ( + os.path.join(test_lib.TEST_DATA_PATH, 'SYSTEM')), + '%UserProfile%\\AppData\\Local\\Microsoft\\Windows\\UsrClass.dat': ( + os.path.join(test_lib.TEST_DATA_PATH, 'UsrClass.dat')), + '%UserProfile%\\NTUSER.DAT': ( + os.path.join(test_lib.TEST_DATA_PATH, 'NTUSER.DAT'))} def Open(self, path, ascii_codepage='cp1252'): """Opens the Windows Registry file specified by the path. @@ -134,13 +158,12 @@ def Open(self, path, ascii_codepage='cp1252'): SkipTest: if the Windows Registry file does not exist and the test should be skipped. """ - if path == '%SystemRoot%\\System32\\config\\SYSTEM': - path = os.path.join(test_lib.TEST_DATA_PATH, 'SYSTEM') - elif path == '%UserProfile%\\NTUSER.DAT': - path = os.path.join(test_lib.TEST_DATA_PATH, 'NTUSER.DAT') + test_file_path = self._MAPPED_PATHS.get(path, None) + if test_file_path is None: + return None return super(TestREGFWinRegistryFileReaderMapped, self).Open( - path, ascii_codepage=ascii_codepage) + test_file_path, ascii_codepage=ascii_codepage) class RegistryTest(test_lib.BaseTestCase): @@ -155,23 +178,18 @@ def testGetCachedFileByPath(self): win_registry = registry.WinRegistry() - # Note that _GetCachedFileByPath expects the key path to be in - # upper case. + # Note that _GetCachedFileByPath expects the key path to be in upper case. key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' - key_path_prefix, registry_file = win_registry._GetCachedFileByPath( - key_path) + + key_path_prefix, registry_file = win_registry._GetCachedFileByPath(key_path) self.assertIsNone(key_path_prefix) self.assertIsNone(registry_file) win_registry = registry.WinRegistry( registry_file_reader=TestREGFWinRegistryFileReader()) + win_registry.OpenAndMapFile(test_path) - registry_file = win_registry._OpenFile(test_path) - key_path_prefix = win_registry.GetRegistryFileMapping(registry_file) - win_registry.MapFile(key_path_prefix, registry_file) - - key_path_prefix, registry_file = win_registry._GetCachedFileByPath( - key_path) + key_path_prefix, registry_file = win_registry._GetCachedFileByPath(key_path) self.assertEqual(key_path_prefix, key_path) self.assertIsNotNone(registry_file) @@ -202,12 +220,7 @@ def testGetKeyByPathFromFile(self): win_registry = registry.WinRegistry( registry_file_reader=TestREGFWinRegistryFileReader()) - - registry_file = win_registry._OpenFile(test_path) - - win_registry = registry.WinRegistry() - key_path_prefix = win_registry.GetRegistryFileMapping(registry_file) - win_registry.MapFile(key_path_prefix, registry_file) + win_registry.OpenAndMapFile(test_path) # Test an existing key. registry_key = win_registry._GetKeyByPathFromFile( @@ -244,9 +257,7 @@ def testGetUsersVirtualKey(self): profile_path = '%SystemRoot%\\System32\\config\\systemprofile\\NTUSER.DAT' win_registry.MapUserFile(profile_path, registry_file) - registry_file = win_registry._OpenFile(software_test_path) - key_path_prefix = win_registry.GetRegistryFileMapping(registry_file) - win_registry.MapFile(key_path_prefix, registry_file) + win_registry.OpenAndMapFile(software_test_path) registry_key = win_registry._GetUsersVirtualKey('\\S-1-5-18') self.assertIsNotNone(registry_key) @@ -276,15 +287,12 @@ def testGetFileByPath(self): test_path = self._GetTestFilePath(['SYSTEM']) self._SkipIfPathNotExists(test_path) - key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' - # Test mapped file with key path prefix. win_registry = registry.WinRegistry( registry_file_reader=TestREGFWinRegistryFileReader()) + win_registry.OpenAndMapFile(test_path) - registry_file = win_registry._OpenFile(test_path) - key_path_prefix = win_registry.GetRegistryFileMapping(registry_file) - win_registry.MapFile(key_path_prefix, registry_file) + key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) self.assertEqual(key_path_prefix, key_path) @@ -297,6 +305,8 @@ def testGetFileByPath(self): registry_file = win_registry._OpenFile(test_path) win_registry.MapFile('', registry_file) + key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) self.assertEqual(key_path_prefix, key_path) self.assertIsNone(registry_file) @@ -304,6 +314,8 @@ def testGetFileByPath(self): # Test without mapped file. win_registry = registry.WinRegistry() + key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) self.assertEqual(key_path_prefix, key_path) self.assertIsNone(registry_file) @@ -312,6 +324,36 @@ def testGetFileByPath(self): win_registry = registry.WinRegistry( registry_file_reader=TestREGFWinRegistryFileReaderMapped()) + key_path = 'HKEY_LOCAL_MACHINE\\SYSTEM' + + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) + self.assertEqual(key_path_prefix, key_path) + self.assertIsNotNone(registry_file) + + key_path = 'HKEY_CURRENT_USER' + + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) + self.assertEqual(key_path_prefix, key_path) + self.assertIsNotNone(registry_file) + + key_path = 'HKEY_CURRENT_USER\\SOFTWARE\\CLASSES' + + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) + self.assertEqual(key_path_prefix, key_path) + self.assertIsNotNone(registry_file) + + # Tests open file based on predefined mapping in reversed order. + win_registry = registry.WinRegistry( + registry_file_reader=TestREGFWinRegistryFileReaderMapped()) + + key_path = 'HKEY_CURRENT_USER\\SOFTWARE\\CLASSES' + + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) + self.assertEqual(key_path_prefix, key_path) + self.assertIsNotNone(registry_file) + + key_path = 'HKEY_CURRENT_USER' + key_path_prefix, registry_file = win_registry._GetFileByPath(key_path) self.assertEqual(key_path_prefix, key_path) self.assertIsNotNone(registry_file) @@ -323,12 +365,7 @@ def testGetKeyByPathOnNTUserDat(self): win_registry = registry.WinRegistry( registry_file_reader=TestREGFWinRegistryFileReader()) - - registry_file = win_registry._OpenFile(test_path) - - win_registry = registry.WinRegistry() - key_path_prefix = win_registry.GetRegistryFileMapping(registry_file) - win_registry.MapFile(key_path_prefix, registry_file) + win_registry.OpenAndMapFile(test_path) # Test an existing key. registry_key = win_registry.GetKeyByPath(