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

chrome/test/mini_installer/registry.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 entry
+
+class RegistryKey(entry.Entry):
+
+  def __init__(self, variable_expander, visitor):
+    super(RegistryKey, self).__init__(variable_expander, visitor)
+    self.key = None
+    self.sub_key = None
+    self.root_key_value = None
+    self.should_key_exist = None
+    self.key_handle = None
+
+  def Check(self, entry_name, entry_value):
+    self.key = self._variable_expander.Expand(entry_name)
+    root_key, self.sub_key = self.key.split('\\', 1)
+    self.root_key_value = self._RootKeyConstant(root_key)
+    self.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.
+      self.key_handle = _winreg.OpenKey(self.root_key_value, self.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.
+      pass
+    # The key exists, see that it matches the expectation.
+    self._visitor.VisitRegistryKey(self)
+
+  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]
+
+
+class RegistryValue(entry.Entry):
+  def __init__(self, variable_expander, visitor):
+    super(RegistryValue, self).__init__(variable_expander, visitor)
+    self.key = None
+    self.value = None
+    self.value_data = None
+    self.value_type = None
+    self.expected_data = None
+    self.expected_type = None
+
+  def Check(self, value, value_expectation, registry_key):
+    # Query the value. It will throw a WindowsError if the value doesn't
+    # exist.
+    self.value = value
+    self.key = registry_key.key
+    self.value_data = None
+    self.value_type = None
+    self.expected_data = value_expectation.get('data', None)
+    if isinstance(self.expected_data, basestring):
+      self.expected_data = self._variable_expander.Expand(self.expected_data)
+    self.expected_type = self._ValueTypeConstant(
+        value_expectation.get('type', None))
+
+    try:
+      self.value_data, self.value_type = _winreg.QueryValueEx(
+          registry_key.key_handle, value)
+    except WindowsError:
+      pass
+
+    self._visitor.VisitRegistryValue(self)
+
+  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 Registry(entry.Entry):
+  """ Visitor for registry entry in property.
+  """
+  def Check(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.
+    """
+    registry_key = RegistryKey(self._variable_expander, self._visitor)
+    registry_key.Check(entry_name, entry_value)
+
+    # Verify the expected values if key exists.
+    if registry_key.key_handle is None or 'values' not in entry_value:
+      return
+
+    for value, value_expectation in entry_value['values'].iteritems():
+      registry_value = RegistryValue(self._variable_expander, self._visitor)
+      registry_value.Check(value, value_expectation, registry_key)