Changes for mapping HKEY_CURRENT_USER\Software\Classes (#222)

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

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,13 +305,17 @@ 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)
 
     # 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(