Home | History | Annotate | Download | only in mini_installer 
      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))
    102