Update policy backend and testserver for the newest policy protocol

net/tools/testserver/device_management.py

 #!/usr/bin/python2.5
 # Copyright (c) 2011 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
@@ skipping 43 matching lines @@
   import simplejson as json
 except ImportError:
   try:
     import json
   except ImportError:
     json = None
 
 import device_management_backend_pb2 as dm
 import cloud_policy_pb2 as cp
 
+# TODO(gfeher): Remove before commit.
+LOG_FILENAME = '/dev/stdout'
+logging.basicConfig(filename=LOG_FILENAME,level=logging.DEBUG)

[gfeher, 2011/02/25 20:36:35]

I didn't forget this!

 
 class RequestHandler(object):
   """Decodes and handles device management requests from clients.
 
   The handler implements all the request parsing and protobuf message decoding
   and encoding. It calls back into the server to lookup, register, and
   unregister clients.
   """
 
   def __init__(self, server, path, headers, request):
@@ skipping 33 matching lines @@
 
     Parses the data supplied at construction time and returns a pair indicating
     http status code and response data to be sent back to the client.
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
     rmsg = dm.DeviceManagementRequest()
     rmsg.ParseFromString(self._request)
 
+    logging.debug('auth -> ' + self._headers.getheader('Authorization', ''))
+    logging.debug('deviceid -> ' + self.GetUniqueParam('deviceid'))
     self.DumpMessage('Request', rmsg)
 
     request_type = self.GetUniqueParam('request')
+    # Check server side requirements, as defined in
+    # device_management_backend.proto.
+    if (self.GetUniqueParam('devicetype') != '2' or
+        self.GetUniqueParam('apptype') != 'Chrome' or
+        (request_type != 'ping' and
+         len(self.GetUniqueParam('deviceid')) >= 64) or
+        len(self.GetUniqueParam('agent')) >= 64):
+      return (400, 'Invalid request parameter')
     if request_type == 'register':
       return self.ProcessRegister(rmsg.register_request)
     elif request_type == 'unregister':
       return self.ProcessUnregister(rmsg.unregister_request)
-    elif request_type == 'policy':
-      return self.ProcessPolicy(rmsg.policy_request)
-    elif request_type == 'cloud_policy':
-      return self.ProcessCloudPolicyRequest(rmsg.cloud_policy_request)
-    elif request_type == 'managed_check':
-      return self.ProcessManagedCheck(rmsg.managed_check_request)
+    elif request_type == 'policy' or request_type == 'ping':
+      return self.ProcessPolicy(rmsg.policy_request, request_type)
     else:
       return (400, 'Invalid request parameter')
 
   def CheckGoogleLogin(self):
     """Extracts the GoogleLogin auth token from the HTTP request, and
     returns it. Returns None if the token is not present.
     """
     match = re.match('GoogleLogin auth=(\\w+)',
                      self._headers.getheader('Authorization', ''))
     if not match:
       return None
     return match.group(1)
 
-  def GetDeviceName(self):
-    """Returns the name for the currently authenticated device based on its
-    device id.
-    """
-    return 'chromeos-' + self.GetUniqueParam('deviceid')
-
   def ProcessRegister(self, msg):
     """Handles a register request.
 
     Checks the query for authorization and device identifier, registers the
     device with the server and constructs a response.
 
     Args:
       msg: The DeviceRegisterRequest message received from the client.
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
     # Check the auth token and device ID.
     if not self.CheckGoogleLogin():
       return (403, 'No authorization')
 
     device_id = self.GetUniqueParam('deviceid')
     if not device_id:
       return (400, 'Missing device identifier')
 
-    # Register the device and create a token.
-    dmtoken = self._server.RegisterDevice(device_id)
+    token_info = self._server.RegisterDevice(device_id,
+                                             msg.machine_id,
+                                             msg.type)
 
     # Send back the reply.
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
-    response.register_response.device_management_token = dmtoken
+    response.register_response.device_management_token = (
+        token_info['device_token'])
+    response.register_response.machine_name = token_info['machine_name']
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
   def ProcessUnregister(self, msg):
     """Handles a register request.
 
     Checks for authorization, unregisters the device and constructs the
     response.
@@ skipping 14 matching lines @@
 
     # Prepare and send the response.
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
     response.unregister_response.CopyFrom(dm.DeviceUnregisterResponse())
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
-  def ProcessManagedCheck(self, msg):
-    """Handles a 'managed check' request.
+  def ProcessInitialPolicy(self, msg):
+    """Handles a 'preregister policy' request.
 
     Queries the list of managed users and responds the client if their user
     is managed or not.
 
     Args:
-      msg: The ManagedCheckRequest message received from the client.
+      msg: The PolicyFetchRequest message received from the client.
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
-    # Check the management token.
+    # Check the GAIA token.
     auth = self.CheckGoogleLogin()
     if not auth:
       return (403, 'No authorization')
 
-    managed_check_response = dm.ManagedCheckResponse()
+    chrome_initial_settings = dm.ChromeInitialSettingsProto()
     if ('*' in self._server.policy['managed_users'] or
         auth in self._server.policy['managed_users']):
-      managed_check_response.mode = dm.ManagedCheckResponse.MANAGED;
+      chrome_initial_settings.enrollment_provision = (
+          dm.ChromeInitialSettingsProto.MANAGED);
     else:
-      managed_check_response.mode = dm.ManagedCheckResponse.UNMANAGED;
+      chrome_initial_settings.enrollment_provision = (
+          dm.ChromeInitialSettingsProto.UNMANAGED);
+
+    policy_data = dm.PolicyData()
+    policy_data.policy_type = msg.policy_type
+    policy_data.policy_value = chrome_initial_settings.SerializeToString()
 
     # Prepare and send the response.
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
-    response.managed_check_response.CopyFrom(managed_check_response)
+    fetch_response = response.policy_response.response.add()
+    fetch_response.policy_data = (
+        policy_data.SerializeToString())
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
-  def ProcessPolicy(self, msg):
-    """Handles a policy request.
-
-    Checks for authorization, encodes the policy into protobuf representation
-    and constructs the response.
-
-    Args:
-      msg: The DevicePolicyRequest message received from the client.
-
-    Returns:
-      A tuple of HTTP status code and response data to send to the client.
-    """
+  def ProcessDevicePolicy(self, msg):
     # Check the management token.
     token, response = self.CheckToken()
     if not token:
       return response
 
     # Stuff the policy dictionary into a response message and send it back.
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
     response.policy_response.CopyFrom(dm.DevicePolicyResponse())
 
     # Respond only if the client requested policy for the cros/device scope,
     # since that's where chrome policy is supposed to live in.
-    if msg.policy_scope in self._server.policy:
-      policy = self._server.policy[msg.policy_scope]['mandatory']
+    if msg.policy_scope == 'chromeos/device':
+      policy = self._server.policy['google/chromeos/user']['mandatory']
       setting = response.policy_response.setting.add()
       setting.policy_key = 'chrome-policy'
       policy_value = dm.GenericSetting()
       for (key, value) in policy.iteritems():
         entry = policy_value.named_value.add()
         entry.name = key
         entry_value = dm.GenericValue()
         if isinstance(value, bool):
           entry_value.value_type = dm.GenericValue.VALUE_TYPE_BOOL
           entry_value.bool_value = value
         elif isinstance(value, int):
           entry_value.value_type = dm.GenericValue.VALUE_TYPE_INT64
           entry_value.int64_value = value
         elif isinstance(value, str) or isinstance(value, unicode):
           entry_value.value_type = dm.GenericValue.VALUE_TYPE_STRING
           entry_value.string_value = value
         elif isinstance(value, list):
           entry_value.value_type = dm.GenericValue.VALUE_TYPE_STRING_ARRAY
           for list_entry in value:
             entry_value.string_array.append(str(list_entry))
         entry.value.CopyFrom(entry_value)
       setting.policy_value.CopyFrom(policy_value)
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
+  def ProcessPolicy(self, msg, request_type):
+    """Handles a policy request.
+
+    Checks for authorization, encodes the policy into protobuf representation
+    and constructs the response.
+
+    Args:
+      msg: The DevicePolicyRequest message received from the client.
+
+    Returns:
+      A tuple of HTTP status code and response data to send to the client.
+    """
+
+    if msg.request:
+      for request in msg.request:
+        if request.policy_type == 'google/chromeos/unregistered_user':
+          if request_type != 'ping':
+            return (400, 'Invalid request type')
+          return self.ProcessInitialPolicy(request)
+        elif (request.policy_type in
+              ('google/chromeos/user', 'google/chromeos/device')):
+          if request_type != 'policy':
+            return (400, 'Invalid request type')
+          return self.ProcessCloudPolicy(request)
+        else:
+          return (400, 'Invalid policy_type')
+    else:
+      return self.ProcessDevicePolicy(msg)
+
   def SetProtobufMessageField(self, group_message, field, field_value):
     '''Sets a field in a protobuf message.
 
     Args:
       group_message: The protobuf message.
       field: The field of the message to set, it shuold be a member of
           group_message.DESCRIPTOR.fields.
       field_value: The value to set.
     '''
-    if field.label == field.LABEL_REPEATED:
+    if field.type == field.TYPE_BOOL:
+      assert type(field_value) == bool
+    elif field.type == field.TYPE_STRING:
+      assert type(field_value) == str
+    elif field.type == field.TYPE_INT64:
+      assert type(field_value) == int
+    elif (field.type == field.TYPE_MESSAGE and
+          field.message_type.name == 'StringList'):
       assert type(field_value) == list
-      assert field.type == field.TYPE_STRING
-      list_field = group_message.__getattribute__(field.name)
+      entries = group_message.__getattribute__(field.name).entries
       for list_item in field_value:
-        list_field.append(list_item)
+        entries.append(list_item)
+      return
     else:
-      # Simple cases:
-      if field.type == field.TYPE_BOOL:
-        assert type(field_value) == bool
-      elif field.type == field.TYPE_STRING:
-        assert type(field_value) == str
-      elif field.type == field.TYPE_INT64:
-        assert type(field_value) == int
-      else:
-        raise Exception('Unknown field type %s' % field.type_name)
-      group_message.__setattr__(field.name, field_value)
+      raise Exception('Unknown field type %s' % field.type)
+    group_message.__setattr__(field.name, field_value)
 
   def GatherPolicySettings(self, settings, policies):
     '''Copies all the policies from a dictionary into a protobuf of type
     CloudPolicySettings.
 
     Args:
       settings: The destination: a CloudPolicySettings protobuf.
       policies: The source: a dictionary containing policies under keys
           'recommended' and 'mandatory'.
     '''
@@ skipping 13 matching lines @@
           group_message.policy_options.mode = cp.PolicyOptions.MANDATORY
           field_value = policies['mandatory'][field.name]
         elif field.name in policies['recommended']:
           field_value = policies['recommended'][field.name]
         if field_value != None:
           got_fields = True
           self.SetProtobufMessageField(group_message, field, field_value)
       if got_fields:
         settings.__getattribute__(group.name).CopyFrom(group_message)
 
-  def ProcessCloudPolicyRequest(self, msg):
+  def ProcessCloudPolicy(self, msg):
     """Handles a cloud policy request. (New protocol for policy requests.)
 
     Checks for authorization, encodes the policy into protobuf representation,
     signs it and constructs the repsonse.
 
     Args:
       msg: The CloudPolicyRequest message received from the client.
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
-    token, response = self.CheckToken()
-    if not token:
-      return response
+
+    token_info, error = self.CheckToken()
+    if not token_info:
+      return error
 
     settings = cp.CloudPolicySettings()
 
-    if msg.policy_scope in self._server.policy:
-      # Respond is only given if the scope is specified in the config file.
+    if (msg.policy_type in token_info['allowed_policy_types'] and
+        msg.policy_type in self._server.policy):
+      # Response is only given if the scope is specified in the config file.
       # Normally 'chromeos/device' and 'chromeos/user' should be accepted.
       self.GatherPolicySettings(settings,
-                                self._server.policy[msg.policy_scope])
+                                self._server.policy[msg.policy_type])
 
-    # Construct response
-    signed_response = dm.SignedCloudPolicyResponse()
-    signed_response.settings.CopyFrom(settings)
-    signed_response.timestamp = int(time.time())
-    signed_response.request_token = token;
-    signed_response.device_name = self.GetDeviceName()
-
-    cloud_response = dm.CloudPolicyResponse()
-    cloud_response.signed_response = signed_response.SerializeToString()
-    signed_data = cloud_response.signed_response
-    cloud_response.signature = (
-        self._server.private_key.hashAndSign(signed_data).tostring())
-    for certificate in self._server.cert_chain:
-      cloud_response.certificate_chain.append(
-          certificate.writeBytes().tostring())
+    policy_data = dm.PolicyData()
+    policy_data.policy_value = settings.SerializeToString()
+    policy_data.policy_type = msg.policy_type
+    policy_data.timestamp = int(time.time() * 1000)
+    policy_data.request_token = token_info['device_token'];
+    policy_data.machine_name = token_info['machine_name']
+    signed_data = policy_data.SerializeToString()
 
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
-    response.cloud_policy_response.CopyFrom(cloud_response)
+    fetch_response = response.policy_response.response.add()
+    fetch_response.policy_data = signed_data
+    fetch_response.policy_data_signature = (
+        self._server.private_key.hashAndSign(signed_data).tostring())
+    for certificate in self._server.cert_chain:
+      fetch_response.certificate_chain.append(
+          certificate.writeBytes().tostring())
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
   def CheckToken(self):
     """Helper for checking whether the client supplied a valid DM token.
 
     Extracts the token from the request and passed to the server in order to
-    look up the client. Returns a pair of token and error response. If the token
-    is None, the error response is a pair of status code and error message.
+    look up the client.
 
     Returns:
-      A pair of DM token and error response. If the token is None, the message
-      will contain the error response to send back.
+      A pair of token information record and error response. If the first
+      element is None, then the second contains an error code to send back to
+      the client. Otherwise the first element is the same structure that is
+      returned by LookupToken().
     """
     error = None
     dmtoken = None
     request_device_id = self.GetUniqueParam('deviceid')
     match = re.match('GoogleDMToken token=(\\w+)',
                      self._headers.getheader('Authorization', ''))
     if match:
       dmtoken = match.group(1)
     if not dmtoken:
       error = dm.DeviceManagementResponse.DEVICE_MANAGEMENT_TOKEN_INVALID
-    elif (not request_device_id or
-          not self._server.LookupDevice(dmtoken) == request_device_id):
-      error = dm.DeviceManagementResponse.DEVICE_NOT_FOUND
     else:
-      return (dmtoken, None)
+      token_info = self._server.LookupToken(dmtoken)
+      if (not token_info or
+          not request_device_id or
+          token_info['device_id'] != request_device_id):
+        error = dm.DeviceManagementResponse.DEVICE_NOT_FOUND
+      else:
+        return (token_info, None)
 
     response = dm.DeviceManagementResponse()
     response.error = error
 
     self.DumpMessage('Response', response)
 
     return (None, (200, response.SerializeToString()))
 
   def DumpMessage(self, label, msg):
     """Helper for logging an ASCII dump of a protobuf message."""
     logging.debug('%s\n%s' % (label, str(msg)))
 
 class TestServer(object):
   """Handles requests and keeps global service state."""
 
   def __init__(self, policy_path, policy_cert_chain):
     """Initializes the server.
 
     Args:
       policy_path: Names the file to read JSON-formatted policy from.
       policy_cert_chain: List of paths to X.509 certificate files of the
           certificate chain used for signing responses.
     """
-    self._registered_devices = {}
+    self._registered_tokens = {}
     self.policy = {}
     if json is None:
       print 'No JSON module, cannot parse policy information'
     else :
       try:
         self.policy = json.loads(open(policy_path).read())
       except IOError:
         print 'Failed to load policy from %s' % policy_path
 
     self.private_key = None
@@ skipping 16 matching lines @@
     Args:
       path: The request path and query parameters received from the client.
       headers: A rfc822.Message-like object containing HTTP headers.
       request: The request data received from the client as a string.
     Returns:
       A pair of HTTP status code and response data to send to the client.
     """
     handler = RequestHandler(self, path, headers, request)
     return handler.HandleRequest()
 
-  def RegisterDevice(self, device_id):
-    """Registers a device and generate a DM token for it.
+  def RegisterDevice(self, device_id, machine_id, type=None):
+    """Registers a device or user and generates a DM token for it.
 
     Args:
       device_id: The device identifier provided by the client.
 
     Returns:
       The newly generated device token for the device.
     """
     dmtoken_chars = []
     while len(dmtoken_chars) < 32:
       dmtoken_chars.append(random.choice('0123456789abcdef'))
     dmtoken = ''.join(dmtoken_chars)
-    self._registered_devices[dmtoken] = device_id
-    return dmtoken
 
-  def LookupDevice(self, dmtoken):
-    """Looks up a device by DMToken.
+    self._registered_tokens[dmtoken] = {
+      'device_id': device_id,
+      'device_token': dmtoken,
+      'allowed_policy_types': {
+        dm.DeviceRegisterRequest.USER: ['google/chromeos/user'],
+        dm.DeviceRegisterRequest.DEVICE: ['google/chromeos/device'],
+        dm.DeviceRegisterRequest.TT: ['google/chromeos/user'],
+      }[type],

[Jakob Kummerow, 2011/02/28 11:06:31]

Not sure if I like this expression. It's concise, yes, but kinda hacky. Also, I
suggest to remove the default value for |type| ("=None"), because this
dictionary lookup would crash in that case.

[gfeher, 2011/02/28 12:21:32]

How about it now?

+      'machine_name': 'chromeos-' + machine_id,
+    }
+    return self._registered_tokens[dmtoken]
+
+  def LookupToken(self, dmtoken):
+    """Looks up a device or a user by DM token.
 
     Args:
       dmtoken: The device management token provided by the client.
 
     Returns:
-      The corresponding device identifier or None if not found.
+      A dictionary with information about a device or user that is registered by
+      dmtoken, or None if the token is not found.
     """
-    return self._registered_devices.get(dmtoken, None)
+    return self._registered_tokens.get(dmtoken, None)
 
   def UnregisterDevice(self, dmtoken):
     """Unregisters a device identified by the given DM token.
 
     Args:
       dmtoken: The device management token provided by the client.
     """
-    if dmtoken in self._registered_devices:
-      del self._registered_devices[dmtoken]
+    if dmtoken in self._registered_tokens.keys():
+      del self._registered_tokens[dmtoken]