New protocol and testserver for the Chrome-DMServer protocol

net/tools/testserver/device_management.py

Index: net/tools/testserver/device_management.py
diff --git a/net/tools/testserver/device_management.py b/net/tools/testserver/device_management.py
index d71522741a950dcc2cf2eab2ca6373a7bf10fd52..183451bb515bbc1bbe1a79437f3324dd5249d88f 100644
--- a/net/tools/testserver/device_management.py
+++ b/net/tools/testserver/device_management.py
@@ -1,5 +1,5 @@
 #!/usr/bin/python2.5
-# Copyright (c) 2010 The Chromium Authors. All rights reserved.
+# 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.
 
@@ -7,21 +7,47 @@
 
 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
-from files in a directory. The files should contain policy definitions in JSON
-format, using the top-level dictionary as a key/value store. The format is
-identical to what the Linux implementation reads from /etc. Here is an example:
+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
+of managed users.
+
+The format of the file is JSON. The root dictionary contains a list under the
+key "managed_users". It contains auth tokens for which the server will claim
+that the user is managed. The token string "*" indicates that all users are
+claimed to be managed. Other keys in the root dictionary identify request
+scopes. Each request scope is described by a dictionary that holds two
+sub-dictionaries: "mandatory" and "recommended". Both these hold the policy
+definitions as key/value stores, their format is identical to what the Linux
+implementation reads from /etc.
+
+Example:
 
 {
-  "HomepageLocation" : "http://www.chromium.org"
+  "chromeos/device": {
+    "mandatory": {
+      "HomepageLocation" : "http://www.chromium.org"
+    },
+    "recommended": {
+      "JavascriptEnabled": false,
+    },
+  },
+  "managed_users": [
+    "secret123456"
+  ]
 }
 
+
 """
 
 import cgi
 import logging
+import os
 import random
 import re
 import sys
+import time
+import tlslite
+import tlslite.api
 
 # The name and availability of the json module varies in python versions.
 try:
@@ -33,6 +59,8 @@ except ImportError:
     json = None
 
 import device_management_backend_pb2 as dm
+import cloud_policy_pb2 as cp
+
 
 class RequestHandler(object):
   """Decodes and handles device management requests from clients.
@@ -95,9 +123,29 @@ class RequestHandler(object):
       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)
     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.
 
@@ -111,11 +159,8 @@ class RequestHandler(object):
       A tuple of HTTP status code and response data to send to the client.
     """
     # Check the auth token and device ID.
-    match = re.match('GoogleLogin auth=(\\w+)',
-                     self._headers.getheader('Authorization', ''))
-    if not match:
+    if not self.CheckGoogleLogin():
       return (403, 'No authorization')
-    auth_token = match.group(1)
 
     device_id = self.GetUniqueParam('deviceid')
     if not device_id:
@@ -128,6 +173,7 @@ class RequestHandler(object):
     response = dm.DeviceManagementResponse()
     response.error = dm.DeviceManagementResponse.SUCCESS
     response.register_response.device_management_token = dmtoken
+    response.register_response.device_name = self.GetDeviceName()
 
     self.DumpMessage('Response', response)
 
@@ -162,11 +208,44 @@ class RequestHandler(object):
 
     return (200, response.SerializeToString())
 
+  def ProcessManagedCheck(self, msg):
+    """Handles a 'managed check' 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.
+
+    Returns:
+      A tuple of HTTP status code and response data to send to the client.
+    """
+    # Check the management token.
+    auth = self.CheckGoogleLogin()
+    if not auth:
+      return (403, 'No authorization')
+
+    managed_check_response = dm.ManagedCheckResponse()
+    if ('*' in self._server.policy['managed_users'] or
+        auth in self._server.policy['managed_users']):
+      managed_check_response.mode = dm.ManagedCheckResponse.MANAGED;
+    else:
+      managed_check_response.mode = dm.ManagedCheckResponse.UNMANAGED;
+
+    # Prepare and send the response.
+    response = dm.DeviceManagementResponse()
+    response.error = dm.DeviceManagementResponse.SUCCESS
+    response.managed_check_response.CopyFrom(managed_check_response)
+
+    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 repsonse.
+    and constructs the response.
 
     Args:
       msg: The DevicePolicyRequest message received from the client.
@@ -214,6 +293,113 @@ class RequestHandler(object):
 
     return (200, response.SerializeToString())
 
+  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:
+      assert type(field_value) == list
+      assert field.type == field.TYPE_STRING
+      list_field = group_message.__getattribute__(field.name)
+      for list_item in field_value:
+        list_field.append(list_item)
+    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)
+
+  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'.
+    '''
+    for group in settings.DESCRIPTOR.fields:
+      # Create protobuf message for group.
+      group_message = eval('cp.' + group.message_type.name + '()')
+      # We assume that this policy group will be recommended, and only switch
+      # it to mandatory if at least one of its members is mandatory.
+      group_message.policy_options.mode = cp.PolicyOptions.RECOMMENDED
+      # Indicates if at least one field was set in |group_message|.
+      got_fields = False
+      # Iterate over fields of the message and feed them from the
+      # policy config file.
+      for field in group_message.DESCRIPTOR.fields:
+        field_value = None
+        if field.name in policies['mandatory']:
+          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):
+    """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
+
+    settings = cp.CloudPolicySettings()
+
+    if msg.policy_scope in self._server.policy:
+      # Respond 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])
+
+    # 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())
+
+    response = dm.DeviceManagementResponse()
+    response.error = dm.DeviceManagementResponse.SUCCESS
+    response.cloud_policy_response.CopyFrom(cloud_response)
+
+    self.DumpMessage('Response', response)
+
+    return (200, response.SerializeToString())
+
   def CheckToken(self):
     """Helper for checking whether the client supplied a valid DM token.
 
@@ -254,11 +440,13 @@ class RequestHandler(object):
 class TestServer(object):
   """Handles requests and keeps global service state."""
 
-  def __init__(self, policy_path):
+  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.policy = {}
@@ -270,6 +458,20 @@ class TestServer(object):
       except IOError:
         print 'Failed to load policy from %s' % policy_path
 
+    self.private_key = None
+    self.cert_chain = []
+    for cert_path in policy_cert_chain:
+      try:
+        cert_text = open(cert_path).read()
+      except IOError:
+        print 'Failed to load certificate from %s' % cert_path
+      certificate = tlslite.api.X509()
+      certificate.parse(cert_text)
+      self.cert_chain.append(certificate)
+      if self.private_key is None:
+        self.private_key = tlslite.api.parsePEMKey(cert_text, private=True)
+        assert self.private_key != None
+
   def HandleRequest(self, path, headers, request):
     """Handles a request.