chrome/test/mini_installer/registry_verifier.py

   1 # Copyright 2013 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import _winreg
   6
   7 import verifier
   8
   9
  10 class RegistryVerifier(verifier.Verifier):
  11   """Verifies that the current registry matches the specified criteria."""
  12
  13   def _RootKeyConstant(self, root_key):
  14     """Converts a root registry key string into a _winreg.HKEY_* constant."""
  15     root_key_mapping = {
  16         'HKEY_CLASSES_ROOT': _winreg.HKEY_CLASSES_ROOT,
  17         'HKEY_CURRENT_USER': _winreg.HKEY_CURRENT_USER,
  18         'HKEY_LOCAL_MACHINE': _winreg.HKEY_LOCAL_MACHINE,
  19         'HKEY_USERS': _winreg.HKEY_USERS,
  20     }
  21     if root_key not in root_key_mapping:
  22       raise KeyError("Unknown root registry key '%s'" % root_key)
  23     return root_key_mapping[root_key]
  24
  25   def _ValueTypeConstant(self, value_type):
  26     """Converts a registry value type string into a _winreg.REG_* constant."""
  27     value_type_mapping = {
  28         'BINARY': _winreg.REG_BINARY,
  29         'DWORD': _winreg.REG_DWORD,
  30         'DWORD_LITTLE_ENDIAN': _winreg.REG_DWORD_LITTLE_ENDIAN,
  31         'DWORD_BIG_ENDIAN': _winreg.REG_DWORD_BIG_ENDIAN,
  32         'EXPAND_SZ': _winreg.REG_EXPAND_SZ,
  33         'LINK': _winreg.REG_LINK,
  34         'MULTI_SZ': _winreg.REG_MULTI_SZ,
  35         'NONE': _winreg.REG_NONE,
  36         'SZ': _winreg.REG_SZ,
  37     }
  38     if value_type not in value_type_mapping:
  39       raise KeyError("Unknown registry value type '%s'" % value_type)
  40     return value_type_mapping[value_type]
  41
  42   def _VerifyExpectation(self, expectation_name, expectation,
  43                          variable_expander):
  44     """Overridden from verifier.Verifier.
  45
  46     Verifies a registry key according to the |expectation|.
  47
  48     Args:
  49       expectation_name: The registry key being verified. It is expanded using
  50           Expand.
  51       expectation: A dictionary with the following keys and values:
  52           'exists' a boolean indicating whether the registry key should exist.
  53           'values' (optional) a dictionary where each key is a registry value
  54               and its associated value is a dictionary with the following key
  55               and values:
  56                   'type' a string indicating the type of the registry value.
  57                   'data' the associated data of the registry value. If it is a
  58                       string, it is expanded using Expand.
  59       variable_expander: A VariableExpander object.
  60     """
  61     key = variable_expander.Expand(expectation_name)
  62     root_key, sub_key = key.split('\\', 1)
  63     try:
  64       # Query the Windows registry for the registry key. It will throw a
  65       # WindowsError if the key doesn't exist.
  66       key_handle = _winreg.OpenKey(self._RootKeyConstant(root_key), sub_key, 0,
  67                                    _winreg.KEY_QUERY_VALUE)
  68     except WindowsError:
  69       # Key doesn't exist. See that it matches the expectation.
  70       assert not expectation['exists'], ('Registry key %s is missing' %
  71                                          key)
  72       return
  73     # The key exists, see that it matches the expectation.
  74     assert expectation['exists'], ('Registry key %s exists' % key)
  75
  76     # Verify the expected values.
  77     if 'values' not in expectation:
  78       return
  79     for value, value_expectation in expectation['values'].iteritems():
  80       # Query the value. It will throw a WindowsError if the value doesn't
  81       # exist.
  82       try:
  83         data, value_type = _winreg.QueryValueEx(key_handle, value)
  84       except WindowsError:
  85         raise KeyError("Value '%s' of registry key %s is missing" % (
  86             value, key))
  87
  88       # Verify the type of the value.
  89       expected_value_type = value_expectation['type']
  90       assert self._ValueTypeConstant(expected_value_type) == value_type, \
  91           "Value '%s' of registry key %s has unexpected type '%s'" % (
  92               value, key, expected_value_type)
  93
  94       # Verify the associated data of the value.
  95       expected_data = value_expectation['data']
  96       if isinstance(expected_data, basestring):
  97         expected_data = variable_expander.Expand(expected_data)
  98       assert expected_data == data, \
  99           ("Value '%s' of registry key %s has unexpected data.\n"
 100            "  Expected: %s\n"
 101            "  Actual: %s" % (value, key, expected_data, data))