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
5import _winreg
6
7import verifier
8
9
10class 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