Fix policy test server key rotation feature

chrome/browser/policy/test/policy_testserver.py

 # Copyright (c) 2012 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.
 
 """A bare-bones test server for testing cloud policy support.
 
 This implements a simple cloud policy test server that can be used to test
 chrome's device management service client. The policy information is read from
 the file named device_management in the server's data directory. It contains
 enforced and recommended policies for the device and user scope, and a list
@@ skipping 863 matching lines @@
           payload = self.CreatePolicyForExternalPolicyData(policy_key)
       else:
         response.error_code = 400
         response.error_message = 'Invalid policy type'
         return
     else:
       response.error_code = 400
       response.error_message = 'Request not allowed for the token used'
       return
 
-    # Sign with 'current_key_index', defaulting to key 0.
-    signing_key = None
-    req_key = None
-    current_key_index = policy.get('current_key_index', 0)
-    nkeys = len(self.server.keys)
-    if (msg.signature_type == dm.PolicyFetchRequest.SHA1_RSA and
-        current_key_index in range(nkeys)):
-      signing_key = self.server.keys[current_key_index]
-      if msg.public_key_version in range(1, nkeys + 1):
-        # requested key exists, use for signing and rotate.
-        req_key = self.server.keys[msg.public_key_version - 1]['private_key']
+    # Determine the current key on the client.
+    client_key_version = None
+    client_key = None
+    if msg.HasField('public_key_version'):
+      client_key_version = msg.public_key_version
+      client_key = self.server.GetKeyByVersion(client_key_version)
+      if client_key is None:
+        response.error_code = 400
+        response.error_message = 'Invalid public key version'
+        return
+
+    # Choose the key for signing the policy.
+    signing_key_version = self.server.GetKeyVersionForSigning(
+        client_key_version)
+    signing_key = self.server.GetKeyByVersion(signing_key_version)
+    assert signing_key is not None
 
     # Fill the policy data protobuf.
     policy_data = dm.PolicyData()
     policy_data.policy_type = msg.policy_type
     policy_data.timestamp = int(time.time() * 1000)
     policy_data.request_token = token_info['device_token']
     policy_data.policy_value = payload
     policy_data.machine_name = token_info['machine_name']
     policy_data.valid_serial_number_missing = (
         token_info['machine_id'] in BAD_MACHINE_IDS)
     policy_data.settings_entity_id = msg.settings_entity_id
     policy_data.service_account_identity = policy.get(
         'service_account_identity',
         'policy_testserver.py-service_account_identity')
     invalidation_source = policy.get('invalidation_source')
     if invalidation_source is not None:
       policy_data.invalidation_source = invalidation_source
     # Since invalidation_name is type bytes in the proto, the Unicode name
     # provided needs to be encoded as ASCII to set the correct byte pattern.
     invalidation_name = policy.get('invalidation_name')
     if invalidation_name is not None:
       policy_data.invalidation_name = invalidation_name.encode('ascii')
 
-    if signing_key:
-      policy_data.public_key_version = current_key_index + 1
+    if msg.signature_type != dm.PolicyFetchRequest.NONE:
+      policy_data.public_key_version = signing_key_version
 
     if username:
       policy_data.username = username
     else:
       # If the correct |username| is unknown, rely on a manually-configured
       # username from the configuration file or use a default.
       policy_data.username = policy.get('policy_user', 'username@example.com')
     policy_data.device_id = token_info['device_id']
 
     # Set affiliation IDs so that user was managed on the device.
     device_affiliation_ids = policy.get('device_affiliation_ids')
     if device_affiliation_ids:
       policy_data.device_affiliation_ids.extend(device_affiliation_ids)
 
     user_affiliation_ids = policy.get('user_affiliation_ids')
     if user_affiliation_ids:
       policy_data.user_affiliation_ids.extend(user_affiliation_ids)
 
-    signed_data = policy_data.SerializeToString()
+    response.policy_data = policy_data.SerializeToString()
 
-    response.policy_data = signed_data
-    if signing_key:
-      response.policy_data_signature = (
-          bytes(signing_key['private_key'].hashAndSign(signed_data)))
-      if msg.public_key_version != current_key_index + 1:
+    # Sign the serialized policy data
+    if msg.signature_type == dm.PolicyFetchRequest.SHA1_RSA:
+      response.policy_data_signature = bytes(
+          signing_key['private_key'].hashAndSign(response.policy_data))
+      if msg.public_key_version != signing_key_version:
         response.new_public_key = signing_key['public_key']
 
         # Set the verification signature appropriate for the policy domain.
         # TODO(atwilson): Use the enrollment domain for public accounts when
         # we add key validation for ChromeOS (http://crbug.com/328038).
         if 'signatures' in signing_key:
           verification_sig = self.GetSignatureForDomain(
               signing_key['signatures'], policy_data.username)
 
           if verification_sig:
             assert len(verification_sig) == 256, \
                 'bad signature size: %d' % len(verification_sig)
             response.new_public_key_verification_signature_deprecated = (
                 verification_sig)
 
-        if req_key:
-          response.new_public_key_signature = (
-              bytes(req_key.hashAndSign(response.new_public_key)))
+        if client_key is not None:
+          response.new_public_key_signature = bytes(
+              client_key['private_key'].hashAndSign(response.new_public_key))
 
     return (200, response.SerializeToString())
 
   def GetSignatureForDomain(self, signatures, username):
     parsed_username = username.split("@", 1)
     if len(parsed_username) != 2:
       logging.error('Could not extract domain from username: %s' % username)
       return None
     domain = parsed_username[1]
 
@@ skipping 45 matching lines @@
   def DumpMessage(self, label, msg):
     """Helper for logging an ASCII dump of a protobuf message."""
     logging.debug('%s\n%s' % (label, str(msg)))
 
 
 class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
                        testserver_base.StoppableHTTPServer):
   """Handles requests and keeps global service state."""
 
   def __init__(self, server_address, data_dir, policy_path, client_state_file,
-               private_key_paths, server_base_url):
+               private_key_paths, rotate_keys_automatically, server_base_url):
     """Initializes the server.
 
     Args:
       server_address: Server host and port.
       policy_path: Names the file to read JSON-formatted policy from.
       private_key_paths: List of paths to read private keys from.
+      rotate_keys_automatically: Whether the keys should be rotated in a
+        round-robin fashion for each policy request (by default, either the
+        key specified in the config or the first key will be used for all

[Andrew T Wilson (Slow), 2016/11/25 14:50:09]

This is fine. As we discussed, we might get slightly more robust test coverage
if we rotate every time by default, but I'm OK with however you want to do it.

[emaxx, 2016/11/25 16:17:38]

This is a valid concern.
I think ideally it should be tackled from this perspective: the existing tests
should test what they were intended for, while the gaps in the test coverage
with regard to key rotation should be filled with new tests.
I've filed a bug http://crbug.com/668716 to track this.

+        requests).
     """
     testserver_base.StoppableHTTPServer.__init__(self, server_address,
                                                  PolicyRequestHandler)
     self._registered_tokens = {}
     self.data_dir = data_dir
     self.policy_path = policy_path
     self.client_state_file = client_state_file
+    self.rotate_keys_automatically = rotate_keys_automatically
     self.server_base_url = server_base_url
 
     self.keys = []
     if private_key_paths:
       # Load specified keys from the filesystem.
       for key_path in private_key_paths:
         try:
           key_str = open(key_path).read()
         except IOError:
           print 'Failed to load private key from %s' % key_path
@@ skipping 58 matching lines @@
     policy = {}
     if json is None:
       logging.error('No JSON module, cannot parse policy information')
     else :
       try:
         policy = json.loads(open(self.policy_path).read(), strict=False)
       except IOError:
         logging.error('Failed to load policies from %s' % self.policy_path)
     return policy
 
+  def GetKeyByVersion(self, key_version):
+    """Obtains the object containing key properties, given the key version.
+
+    Args:
+      key_version: Integer key version.
+
+    Returns:
+      The object containing key properties, or None if the key is not found.
+    """
+    key_index = key_version - 1
+    if key_index < 0:
+      return None
+    if key_index >= len(self.keys):
+      if self.rotate_keys_automatically:
+        key_index %= len(self.keys)
+      else:
+        return None
+    return self.keys[key_index]
+
+  def GetKeyVersionForSigning(self, client_key_version):
+    """Determines the version of the key that should be used for signing policy.
+
+    Args:
+      client_key_version: Either an integer representing the current key version
+        provided by the client, or None if the client didn't provide any.
+
+    Returns:
+      An integer representing the signing key version.
+    """
+    if self.rotate_keys_automatically and client_key_version is not None:
+      return client_key_version + 1
+    return self.GetPolicies().get('current_key_index', 0) + 1
+
   def ResolveUser(self, auth_token):
     """Tries to resolve an auth token to the corresponding user name.
 
     If enabled, this makes a request to the token info endpoint to determine the
     user ID corresponding to the token. If token resolution is disabled or the
     request fails, this will return the policy_user config parameter.
     """
     config = self.GetPolicies()
     token_info_url = config.get('token_info_url')
     if token_info_url is not None:
@@ skipping 245 matching lines @@
     super(PolicyServerRunner, self).__init__()
 
   def create_server(self, server_data):
     data_dir = self.options.data_dir or ''
     config_file = (self.options.config_file or
                    os.path.join(data_dir, 'device_management'))
     server = PolicyTestServer((self.options.host, self.options.port),
                               data_dir, config_file,
                               self.options.client_state_file,
                               self.options.policy_keys,
+                              self.options.rotate_keys_automatically,
                               self.options.server_base_url)
     server_data['port'] = server.server_port
     return server
 
   def add_options(self):
     testserver_base.TestServerRunner.add_options(self)
     self.option_parser.add_option('--client-state', dest='client_state_file',
                                   help='File that client state should be '
                                   'persisted to. This allows the server to be '
                                   'seeded by a list of pre-registered clients '
                                   'and restarts without abandoning registered '
                                   'clients.')
     self.option_parser.add_option('--policy-key', action='append',
                                   dest='policy_keys',
                                   help='Specify a path to a PEM-encoded '
                                   'private key to use for policy signing. May '
                                   'be specified multiple times in order to '
-                                  'load multiple keys into the server. If the '
-                                  'server has multiple keys, it will rotate '
-                                  'through them in at each request in a '
-                                  'round-robin fashion. The server will '
-                                  'use a canned key if none is specified '
-                                  'on the command line. The test server will '
-                                  'also look for a verification signature file '
-                                  'in the same location: <filename>.sig and if '
-                                  'present will add the signature to the '
-                                  'policy blob as appropriate via the '
+                                  'load multiple keys into the server. The '
+                                  'server will use a canned key if none is '
+                                  'specified on the command line. The test '
+                                  'server will also look for a verification '
+                                  'signature file in the same location: '
+                                  '<filename>.sig and if present will add the '
+                                  'signature to the policy blob as appropriate '
+                                  'via the '
                              'new_public_key_verification_signature_deprecated '
                                   'field.')
+    self.option_parser.add_option('--rotate-policy-keys-automatically',
+                                  action='store_true',
+                                  dest='rotate_keys_automatically',
+                                  help='If present, then the policy keys will '
+                                  'be rotated in a round-robin fashion for '
+                                  'each policy request (by default, either the '
+                                  'key specified in the config or the first '
+                                  'key will be used for all requests).')
     self.option_parser.add_option('--log-level', dest='log_level',
                                   default='WARN',
                                   help='Log level threshold to use.')
     self.option_parser.add_option('--config-file', dest='config_file',
                                   help='Specify a configuration file to use '
                                   'instead of the default '
                                   '<data_dir>/device_management')
     self.option_parser.add_option('--server-base-url', dest='server_base_url',
                                   help='The server base URL to use when '
                                   'constructing URLs to return to the client.')
 
   def run_server(self):
     logger = logging.getLogger()
     logger.setLevel(getattr(logging, str(self.options.log_level).upper()))
     if (self.options.log_to_console):
       logger.addHandler(logging.StreamHandler())
     if (self.options.log_file):
       logger.addHandler(logging.FileHandler(self.options.log_file))
 
     testserver_base.TestServerRunner.run_server(self)
 
 
 if __name__ == '__main__':
   sys.exit(PolicyServerRunner().main())