Changes to REGFWinRegistryFile for virtual keys

dfwinreg/creg.py

@@ -39,7 +39,7 @@ def GetKeyByPath(self, key_path):
       key_path (str): Windows Registry key path.
 
     Returns:
-      WinRegistryKey: Registry key or None if not available.
+      WinRegistryKey: Windows Registry key or None if not available.
     """
     key_path_upper = key_path.upper()
     if key_path_upper.startswith(self._key_path_prefix_upper):

dfwinreg/regf.py

@@ -16,7 +16,7 @@ class REGFWinRegistryFile(interface.WinRegistryFile):
   """Implementation of a Windows Registry file using pyregf."""
 
   def __init__(
-      self, ascii_codepage='cp1252', emulate_virtual_keys=True,
+      self, ascii_codepage='cp1252', emulate_virtual_keys=False,
       key_path_prefix=''):
     """Initializes the Windows Registry file.
 
@@ -34,12 +34,11 @@ def __init__(
     self._regf_file = pyregf.file()
     self._regf_file.set_ascii_codepage(ascii_codepage)
 
-  def _GetCurrentControlSetKeyPath(self):
-    """Retrieves the key path of the current control set key.
+  def _GetCurrentControlSetKey(self):
+    """Retrieves the current control set key.
 
     Returns:
-      str: key path of the current control set Windows Registry key or None
-          if not available.
+      pyregf.key: current control key or None if not available.
     """
     select_key = self._GetKeyByPathFromFile('\\Select')
     if not select_key:
@@ -63,7 +62,7 @@ def _GetCurrentControlSetKeyPath(self):
     if not control_set or control_set <= 0 or control_set > 999:
       return None
 
-    return f'\\ControlSet{control_set:03d}'
+    return self._GetKeyByPathFromFile(f'\\ControlSet{control_set:03d}')
 
   def _GetKeyByPathFromFile(self, relative_key_path):
     """Retrieves the key for a specific path form the Windows Registry file.
@@ -72,13 +71,24 @@ def _GetKeyByPathFromFile(self, relative_key_path):
       relative_key_path (str): Windows Registry key path relative to the file.
 
     Returns:
-      pyregf.key: Registry key or None if not available.
+      pyregf.key: Windows Registry key or None if not available.
     """
     try:
       return self._regf_file.get_key_by_path(relative_key_path)
     except IOError:
       return None
 
+  def AddCurrentControlSetKey(self):
+    """Adds a virtual current control set key.
+
+    Raises:
+      ValueError: if the virtual key already exists.
+    """
+    pyregf_key = self._GetCurrentControlSetKey()
+    if pyregf_key:
+      self._key_helper.AddVirtualKey('\\CurrentControlSet', pyregf_key)
+      self._emulate_virtual_keys = True
+
   def AddVirtualKey(self, relative_key_path, pyregf_key):
     """Adds a virtual key.
 
@@ -90,6 +100,7 @@ def AddVirtualKey(self, relative_key_path, pyregf_key):
       ValueError: if the virtual key already exists.
     """
     self._key_helper.AddVirtualKey(relative_key_path, pyregf_key)
+    self._emulate_virtual_keys = True
 
   def Close(self):
     """Closes the Windows Registry file."""
@@ -103,7 +114,7 @@ def GetKeyByPath(self, key_path):
       key_path (str): Windows Registry key path.
 
     Returns:
-      WinRegistryKey: Registry key or None if not available.
+      WinRegistryKey: Windows Registry key or None if not available.
     """
     key_path_upper = key_path.upper()
     if key_path_upper.startswith(self._key_path_prefix_upper):
@@ -151,13 +162,6 @@ def Open(self, file_object):
     self._file_object = file_object
     self._regf_file.open_file_object(self._file_object)
 
-    if self._emulate_virtual_keys:
-      key_path = self._GetCurrentControlSetKeyPath()
-      if key_path:
-        pyregf_key = self._GetKeyByPathFromFile(key_path)
-        if pyregf_key:
-          self.AddVirtualKey('\\CurrentControlSet', pyregf_key)
-
     return True
 
 

dfwinreg/registry.py

@@ -286,16 +286,6 @@ def _GetFileByPath(self, key_path_upper):
     key_path_prefix_upper = key_path_prefix.upper()
     self.MapFile(key_path_prefix_upper, registry_file)
 
-    # If HKEY_CURRENT_USER set HKEY_CURRENT_USER\\Software\\Classes as a 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 and isinstance(registry_key, regf.REGFWinRegistryKey):
-        # 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):
@@ -583,9 +573,27 @@ def MapFile(self, key_path_prefix, registry_file):
       key_path_prefix (str): key path prefix.
       registry_file (WinRegistryFile): Windows Registry file.
     """
-    self._registry_files[key_path_prefix.upper()] = registry_file
+    key_path_prefix_upper = key_path_prefix.upper()
+
+    self._registry_files[key_path_prefix_upper] = registry_file
     registry_file.SetKeyPathPrefix(key_path_prefix)
 
+    # If HKEY_CURRENT_USER set HKEY_CURRENT_USER\\Software\\Classes as a 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 and isinstance(registry_key, regf.REGFWinRegistryKey):
+        # pylint: disable=protected-access
+        pyregf_key = registry_key._pyregf_key
+        registry_file.AddVirtualKey('\\Software\\Classes', pyregf_key)
+
+    # If HKEY_LOCAL_MACHINE\\System set HKEY_LOCAL_MACHINE\\System\\
+    # CurrentControlSet as a virtual key in the file.
+    elif key_path_prefix_upper == 'HKEY_LOCAL_MACHINE\\SYSTEM' and isinstance(
+        registry_file, regf.REGFWinRegistryFile):
+      registry_file.AddCurrentControlSetKey()
+
   def MapUserFile(self, profile_path, registry_file):
     """Maps the user Windows Registry file to a specific profile path.
 

tests/regf.py

@@ -113,6 +113,7 @@ def testGetRootKey(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         registry_key = registry_file.GetRootKey()
@@ -148,6 +149,7 @@ def testGetKeyByPath(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         key_path = '\\'
@@ -465,6 +467,7 @@ def testGetSubkeyByIndex(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         key_path = '\\'
@@ -502,6 +505,7 @@ def testGetSubkeyByName(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         key_path = '\\'
@@ -542,6 +546,7 @@ def testGetSubkeyByPath(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         key_path = '\\'
@@ -588,6 +593,7 @@ def testGetSubkeys(self):
 
     with open(test_path, 'rb') as file_object:
       registry_file.Open(file_object)
+      registry_file.AddCurrentControlSetKey()
 
       try:
         key_path = '\\'

tests/registry.py

@@ -31,8 +31,8 @@ def GetKeyByPath(self, key_path):
     """
     if key_path == 'HKEY_LOCAL_MACHINE\\System\\Select':
       registry_key = fake.FakeWinRegistryKey(
-          'Select', key_path='HKEY_LOCAL_MACHINE\\System',
-          last_written_time=0)
+          'Select', key_path_prefix='HKEY_LOCAL_MACHINE\\System',
+          last_written_time=0, relative_key_path='')
 
       registry_value = fake.FakeWinRegistryValue(
           'Current', data=b'DATA', data_type=definitions.REG_BINARY)
@@ -78,15 +78,9 @@ def _GetFileByPath(self, key_path_upper):
 class TestREGFWinRegistryFileReader(interface.WinRegistryFileReader):
   """Windows Registry file reader that reads a single test file."""
 
-  def __init__(self, emulate_virtual_keys=True):
-    """Initializes a file Windows Registry file reader.
-
-    Args:
-      emulate_virtual_keys (Optional[bool]): True if virtual keys should be
-          emulated.
-    """
+  def __init__(self):
+    """Initializes a file Windows Registry file reader."""
     super(TestREGFWinRegistryFileReader, self).__init__()
-    self._emulate_virtual_keys = emulate_virtual_keys
     self._file_objects = []
 
   def __del__(self):
@@ -108,9 +102,7 @@ def Open(self, path, ascii_codepage='cp1252'):
     if not os.path.exists(path):
       return None
 
-    registry_file = regf.REGFWinRegistryFile(
-        ascii_codepage=ascii_codepage,
-        emulate_virtual_keys=self._emulate_virtual_keys)
+    registry_file = regf.REGFWinRegistryFile(ascii_codepage=ascii_codepage)
 
     file_object = open(path, 'rb')  # pylint: disable=consider-using-with
 

tests/registry_searcher.py

@@ -269,11 +269,7 @@ def testFind(self):
 
     win_registry = registry.WinRegistry(
         registry_file_reader=test_registry.TestREGFWinRegistryFileReader())
-
-    registry_file = win_registry._OpenFile(test_path)
-
-    key_path_prefix = win_registry.GetRegistryFileMapping(registry_file)
-    win_registry.MapFile(key_path_prefix, registry_file)
+    win_registry.OpenAndMapFile(test_path)
 
     searcher = registry_searcher.WinRegistrySearcher(win_registry)
 
@@ -326,17 +322,12 @@ def testFind(self):
 
     # Test without find specifications.
     win_registry = registry.WinRegistry(
-        registry_file_reader=test_registry.TestREGFWinRegistryFileReader(
-            emulate_virtual_keys=False))
-
-    registry_file = win_registry._OpenFile(test_path)
-
-    key_path_prefix = win_registry.GetRegistryFileMapping(registry_file)
-    win_registry.MapFile(key_path_prefix, registry_file)
+        registry_file_reader=test_registry.TestREGFWinRegistryFileReader())
+    win_registry.OpenAndMapFile(test_path)
 
     searcher = registry_searcher.WinRegistrySearcher(win_registry)
     key_paths = list(searcher.Find())
-    self.assertEqual(len(key_paths), 21512)
+    self.assertEqual(len(key_paths), 31353)
 
   def testGetKeyByPath(self):
     """Tests the GetKeyByPath function."""
@@ -345,11 +336,7 @@ def testGetKeyByPath(self):
 
     win_registry = registry.WinRegistry(
         registry_file_reader=test_registry.TestREGFWinRegistryFileReader())
-
-    registry_file = win_registry._OpenFile(test_path)
-
-    key_path_prefix = win_registry.GetRegistryFileMapping(registry_file)
-    win_registry.MapFile(key_path_prefix, registry_file)
+    win_registry.OpenAndMapFile(test_path)
 
     searcher = registry_searcher.WinRegistrySearcher(win_registry)
 
@@ -368,11 +355,7 @@ def testSplitKeyPath(self):
 
     win_registry = registry.WinRegistry(
         registry_file_reader=test_registry.TestREGFWinRegistryFileReader())
-
-    registry_file = win_registry._OpenFile(test_path)
-
-    key_path_prefix = win_registry.GetRegistryFileMapping(registry_file)
-    win_registry.MapFile(key_path_prefix, registry_file)
+    win_registry.OpenAndMapFile(test_path)
 
     searcher = registry_searcher.WinRegistrySearcher(win_registry)