Cleanup machine based on the state in configuration file for mini installer test.

chrome/test/mini_installer/registry_visitor.py

+# Copyright 2017 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import _winreg
+
+import visitor
+
+class RegistryVisitor(visitor.Visitor):
+  """ Visitor for registry entry in property.
+  """
+  def _VisitEntry(self, entry_name, entry_value):
+    """ Visit one registry entry's key and values. Values are not checked if
+        the key is not present in the registry.
+
+        Args:
+          entry_name: A string represents regsitry name in the property.
+          entry_value: A dictionary represents the following keys and values:
+            'exists' - a string indicating whether the registry key's existence
+                       is 'required', 'optional', or 'forbidden'.
+            'values' - (optional) a dictionary where each key is a registry
+                       value and its associated value is a dictionary with the
+                       following key and values:
+
+                      'type' - (optional) a string indicating the type of the
+                               registry value. If not present, the corresponding
+                               value is expected to be absent in the registry.
+                      'data' - the associated data of the registry value if
+                               'type' is specified. If it is a string, it is
+                               expanded using Expand.
+    """
+    key = self._variable_expander.Expand(entry_name)
+    root_key, sub_key = key.split('\\', 1)
+    root_key_value = self._RootKeyConstant(root_key)
+    should_key_exist = entry_value['exists']
+    try:
+      # Query the Windows registry for the registry key. It will throw a
+      # WindowsError if the key doesn't exist.
+      key_handle = _winreg.OpenKey(root_key_value, sub_key, 0,
+                                   _winreg.KEY_QUERY_VALUE)
+    except WindowsError:
+      # Key doesn't exist. See that it matches the expectation.
+      # Values are not checked if the missing key's existence is optional.
+      return self._HandleRegistryKey(root_key_value, sub_key, False,
+          should_key_exist, key)
+    # The key exists, see that it matches the expectation.
+    self._HandleRegistryKey(root_key_value, sub_key, True, should_key_exist,
+        key)
+
+    # Verify the expected values.
+    if 'values' not in entry_value:
+      return
+    for value, value_expectation in entry_value['values'].iteritems():
+      # Query the value. It will throw a WindowsError if the value doesn't
+      # exist.
+      value_data = None
+      value_type = None
+      expected_data = value_expectation.get('data', None)
+      if isinstance(expected_data, basestring):
+        expected_data = self._variable_expander.Expand(expected_data)
+      expected_type = self._ValueTypeConstant(
+          value_expectation.get('type', None))
+      try:
+        value_data, value_type = _winreg.QueryValueEx(key_handle, value)
+      except WindowsError:
+        pass
+
+      self._HandleRegistryValue(value_data, value_type, expected_data,
+          expected_type, key, value)
+
+
+
+  def _HandleRegistryKey(self, root, key_path, is_key_exists, should_key_exist,
+                         key):
+    """ Abstract function to check a regsitry key.
+
+        Args:
+          root: An integer represnets HKEY_* constant.
+          key_path: A string represents the path of key.
+          is_key_exists: A boolean represnets whether the key exists or not.
+          should_key_exist: a string represents whether the registry key's
+                            existence is 'required', 'optional', or 'forbidden'.
+          key: A string represents the whole name of the key.
+    """
+    raise NotImplementedError()
+
+  def _HandleRegistryValue(self, value_data, value_type, expected_data,
+                           expected_type, key, value):
+    """ Abstract function to check one value of the registry key.
+
+        Args:
+          value_data: A instance represents the data of the registry value.
+          value_type: A _winreg const represents the type of the registry value.
+          expected_data: A instance represents the data of expected registry
+                         value.
+          expected_type: A _winreg const represents the type of expected
+                         registry value.
+          key: A string represents the whole name of the key.
+          value: A string represents the name of the value.
+    """
+    raise NotImplementedError()
+
+  def _RootKeyConstant(self, root_key):
+    """Converts a root registry key string into a _winreg.HKEY_* constant."""
+    root_key_mapping = {
+        'HKEY_CLASSES_ROOT': _winreg.HKEY_CLASSES_ROOT,
+        'HKEY_CURRENT_USER': _winreg.HKEY_CURRENT_USER,
+        'HKEY_LOCAL_MACHINE': _winreg.HKEY_LOCAL_MACHINE,
+        'HKEY_USERS': _winreg.HKEY_USERS,
+    }
+    if root_key not in root_key_mapping:
+      raise KeyError("Unknown root registry key '%s'" % root_key)
+    return root_key_mapping[root_key]
+
+  def _ValueTypeConstant(self, value_type):
+    """Converts a registry value type string into a _winreg.REG_* constant."""
+    if value_type is None:
+      return None
+    value_type_mapping = {
+        'BINARY': _winreg.REG_BINARY,
+        'DWORD': _winreg.REG_DWORD,
+        'DWORD_LITTLE_ENDIAN': _winreg.REG_DWORD_LITTLE_ENDIAN,
+        'DWORD_BIG_ENDIAN': _winreg.REG_DWORD_BIG_ENDIAN,
+        'EXPAND_SZ': _winreg.REG_EXPAND_SZ,
+        'LINK': _winreg.REG_LINK,
+        'MULTI_SZ': _winreg.REG_MULTI_SZ,
+        'NONE': _winreg.REG_NONE,
+        'SZ': _winreg.REG_SZ,
+    }
+    if value_type not in value_type_mapping:
+      raise KeyError("Unknown registry value type '%s'" % value_type)
+    return value_type_mapping[value_type]
+
+class RegistryVerifier(RegistryVisitor):
+  """Verifies that the current registry matches the specified criteria."""
+
+  def _HandleRegistryKey(self, root, key_path, is_key_exists, should_key_exist,
+                         key):
+    """ Overridden RegistryVisitor._HandleRegistryKey to verify a registry key
+        according to the expected key.
+    """
+    if should_key_exist == 'forbidden':
+      assert not is_key_exists, 'Registry key %s exists' % key
+    elif should_key_exist == 'required':
+      assert is_key_exists, 'Registry key %s is missing' % key
+
+  def _HandleRegistryValue(self, value_data, value_type, expected_data,
+                           expected_type, key, value):
+    """ Overridden RegistryVisitor._HandleRegistryValue to verify a registry
+        value according to the expected value.
+    """
+    error_message = ""
+    if expected_type == None:
+      error_message = ('Value %s of registry key %s exists '
+                     'with value %s') % (value, key, value_data)
+    elif value_type == None:
+      error_message = 'Value %s of registry key %s is missing' % (value, key)
+    else:
+      error_message = 'Value %s of registry key %s has unexpected type %s' % (
+              value, key, expected_type)
+    assert expected_type == value_type, error_message
+    assert expected_data == value_data, (
+          'Value %s of registry key %s has unexpected data.\n'
+           '  Expected: %s\n'
+           '  Actual: %s') % (value, key, expected_data, value_data)
+
+class RegistryCleaner(RegistryVisitor):
+  """ Reset the currect registry to match the expectation.
+  """
+  def _HandleRegistryKey(self, root, key_path, is_key_exists, should_key_exist,
+                         key):
+    """ Overridden RegistryVisitor._HandleRegistryKey to delete a registry key
+        if necessary.
+    """
+    if is_key_exists and should_key_exist == 'forbidden':
+      try:
+        self._DeleteRegKey(root, key_path)
+      except WindowsError:
+        # Suppress any delete error
+        pass
+
+  def _HandleRegistryValue(self, value_data, value_type, expected_data,
+                           expected_type, key, value):
+    """ Overridden RegistryVisitor._HandleRegistryValue."""
+    pass
+
+  def _DeleteRegKey(self, root, key_path):
+    """ Delete a registry key recursively.
+
+        Args:
+          root: An integer represents HKEY_ constant.
+          key_path: A string represents the path of key that needs to be
+                    deleted.
+    """
+    with _winreg.OpenKey(root, key_path, 0, _winreg.KEY_SET_VALUE |
+                                            _winreg.KEY_READ |
+                                            _winreg.KEY_WOW64_32KEY) as key:
+      num_of_sub_keys, _, _ =  _winreg.QueryInfoKey(key)
+      for i in range(0, num_of_sub_keys):
+        self._DeleteRegKey(root, key_path + "\\" + _winreg.EnumKey(key, i))
+    _winreg.DeleteKey(root, key_path)