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
+        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.
+    """
+    # Convert the policy key version, which has to be positive according to the
+    # policy protocol definition, to an index in the keys list.
+    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 the incremented version, which means that the key should be
+      # rotated.
+      return client_key_version + 1
+    # Return the version that is specified by the config, defaulting to using
+    # the very first key. Note that incrementing here is done due to conversion
+    # between indices in the keys list and the key versions transmitted to the
+    # client (where the latter have to be positive according to the policy
+    # protocol definition).
+    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())