Split out policy code from net/tools/testserver.

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 33 matching lines @@
     }
   },
   "managed_users" : [
     "secret123456"
   ],
   "current_key_index": 0
 }
 
 """
 
+import BaseHTTPServer
 import cgi
 import hashlib
 import logging
 import os
 import random
 import re
 import sys
 import time
 import tlslite
 import tlslite.api
 import tlslite.utils
 
 # The name and availability of the json module varies in python versions.
 try:
   import simplejson as json
 except ImportError:
   try:
     import json
   except ImportError:
     json = None
 
 import asn1der
+import testserver_base
+
 import device_management_backend_pb2 as dm
 import cloud_policy_pb2 as cp
 import chrome_device_policy_pb2 as dp
 
 # ASN.1 object identifier for PKCS#1/RSA.
 PKCS1_RSA_OID = '\x2a\x86\x48\x86\xf7\x0d\x01\x01\x01'
 
 # SHA256 sum of "0".
 SHA256_0 = hashlib.sha256('0').digest()
 
 # List of bad machine identifiers that trigger the |valid_serial_number_missing|
 # flag to be set set in the policy fetch response.
 BAD_MACHINE_IDS = [ '123490EN400015' ];
 
 # List of machines that trigger the server to send kiosk enrollment response
 # for the register request.
 KIOSK_MACHINE_IDS = [ 'KIOSK' ];
 
-class RequestHandler(object):
+
+class PolicyRequestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
   """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):
+  def __init__(self, request, client_address, server):
     """Initialize the handler.
 
     Args:
+      request: The request data received from the client as a string.
+      client_address: The client address.
       server: The TestServer object to use for (un)registering clients.
-      path: A string containing the request path and query parameters.
-      headers: A rfc822.Message-like object containing HTTP headers.
-      request: The request data received from the client as a string.
     """
-    self._server = server
-    self._path = path
-    self._headers = headers
-    self._request = request
-    self._params = None
+    BaseHTTPServer.BaseHTTPRequestHandler.__init__(self, request,
+                                                   client_address, server)
 
   def GetUniqueParam(self, name):
     """Extracts a unique query parameter from the request.
 
     Args:
       name: Names the parameter to fetch.
     Returns:
       The parameter value or None if the parameter doesn't exist or is not
       unique.
     """
-    if not self._params:
-      self._params = cgi.parse_qs(self._path[self._path.find('?') + 1:])
+    if not hasattr(self, '_params'):
+      self._params = cgi.parse_qs(self.path[self.path.find('?') + 1:])
 
     param_list = self._params.get(name, [])
     if len(param_list) == 1:
       return param_list[0]
     return None;
 
+  def do_POST(self):
+    http_response, raw_reply = self.HandleRequest();
+    self.send_response(http_response)
+    if (http_response == 200):
+      self.send_header('Content-Type', 'application/x-protobuffer')
+    self.end_headers()
+    self.wfile.write(raw_reply)
+
   def HandleRequest(self):
     """Handles a request.
 
     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)
+    length = int(self.headers.getheader('content-length'))
+    rmsg.ParseFromString(self.rfile.read(length))
 
     logging.debug('gaia auth token -> ' +
-                  self._headers.getheader('Authorization', ''))
+                  self.headers.getheader('Authorization', ''))
     logging.debug('oauth token -> ' + str(self.GetUniqueParam('oauth_token')))
     logging.debug('deviceid -> ' + str(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
@@ skipping 15 matching lines @@
     """Extracts the auth token from the request and returns it. The token may
     either be a GoogleLogin token from an Authorization header, or an OAuth V2
     token from the oauth_token query parameter. Returns None if no token is
     present.
     """
     oauth_token = self.GetUniqueParam('oauth_token')
     if oauth_token:
       return oauth_token
 
     match = re.match('GoogleLogin auth=(\\w+)',
-                     self._headers.getheader('Authorization', ''))
+                     self.headers.getheader('Authorization', ''))
     if match:
       return match.group(1)
 
     return None
 
   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.
     auth = self.CheckGoogleLogin()
     if not auth:
       return (403, 'No authorization')
 
-    policy = self._server.GetPolicies()
+    policy = self.server.GetPolicies()
     if ('*' not in policy['managed_users'] and
         auth not in policy['managed_users']):
       return (403, 'Unmanaged')
 
     device_id = self.GetUniqueParam('deviceid')
     if not device_id:
       return (400, 'Missing device identifier')
 
-    token_info = 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.register_response.device_management_token = (
         token_info['device_token'])
     response.register_response.machine_name = token_info['machine_name']
     response.register_response.enrollment_type = token_info['enrollment_mode']
 
@@ skipping 12 matching lines @@
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
     # Check the management token.
     token, response = self.CheckToken();
     if not token:
       return response
 
     # Unregister the device.
-    self._server.UnregisterDevice(token['device_token']);
+    self.server.UnregisterDevice(token['device_token']);
 
     # Prepare and send the response.
     response = dm.DeviceManagementResponse()
     response.unregister_response.CopyFrom(dm.DeviceUnregisterResponse())
 
     self.DumpMessage('Response', response)
 
     return (200, response.SerializeToString())
 
   def ProcessPolicy(self, msg, request_type):
@@ skipping 170 matching lines @@
 
     Returns:
       A tuple of HTTP status code and response data to send to the client.
     """
 
     token_info, error = self.CheckToken()
     if not token_info:
       return error
 
     if msg.machine_id:
-      self._server.UpdateMachineId(token_info['device_token'], msg.machine_id)
+      self.server.UpdateMachineId(token_info['device_token'], msg.machine_id)
 
     # Response is only given if the scope is specified in the config file.
     # Normally 'google/chromeos/device', 'google/chromeos/user' and
     # 'google/chromeos/publicaccount' should be accepted.
-    policy = self._server.GetPolicies()
+    policy = self.server.GetPolicies()
     policy_value = ''
     policy_key = msg.policy_type
     if msg.settings_entity_id:
       policy_key += '/' + msg.settings_entity_id
     if msg.policy_type in token_info['allowed_policy_types']:
       if (msg.policy_type == 'google/chromeos/user' or
           msg.policy_type == 'google/chrome/user' or
           msg.policy_type == 'google/chromeos/publicaccount'):
         settings = cp.CloudPolicySettings()
         self.GatherUserPolicySettings(settings, policy.get(policy_key, {}))
       elif msg.policy_type == 'google/chromeos/device':
         settings = dp.ChromeDeviceSettingsProto()
         self.GatherDevicePolicySettings(settings, policy.get(policy_key, {}))
 
     # 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)
+    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]
+      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']
+        req_key = self.server.keys[msg.public_key_version - 1]['private_key']
 
     # 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 = settings.SerializeToString()
     policy_data.machine_name = token_info['machine_name']
     policy_data.valid_serial_number_missing = (
         token_info['machine_id'] in BAD_MACHINE_IDS)
@@ skipping 37 matching lines @@
     Returns:
       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 = 500
     dmtoken = None
     request_device_id = self.GetUniqueParam('deviceid')
     match = re.match('GoogleDMToken token=(\\w+)',
-                     self._headers.getheader('Authorization', ''))
+                     self.headers.getheader('Authorization', ''))
     if match:
       dmtoken = match.group(1)
     if not dmtoken:
       error = 401
     else:
-      token_info = self._server.LookupToken(dmtoken)
+      token_info = self.server.LookupToken(dmtoken)
       if (not token_info or
           not request_device_id or
           token_info['device_id'] != request_device_id):
         error = 410
       else:
         return (token_info, None)
 
     logging.debug('Token check failed with error %d' % error)
 
     return (None, (error, 'Server error %d' % error))
 
   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):
+
+class PolicyTestServer(testserver_base.ClientRestrictingServerMixIn,
+                       testserver_base.BrokenPipeHandlerMixIn,
+                       testserver_base.StoppableHTTPServer):
   """Handles requests and keeps global service state."""
 
-  def __init__(self, policy_path, private_key_paths):
+  def __init__(self, server_address, policy_path, private_key_paths):
     """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.
     """
+    testserver_base.StoppableHTTPServer.__init__(self, server_address,
+                                                 PolicyRequestHandler)
     self._registered_tokens = {}
     self.policy_path = policy_path
 
     self.keys = []
     if private_key_paths:
       # Load specified keys from the filesystem.
       for key_path in private_key_paths:
         try:
           key = tlslite.api.parsePEMKey(open(key_path).read(), private=True)
         except IOError:
@@ skipping 28 matching lines @@
     policy = {}
     if json is None:
       print 'No JSON module, cannot parse policy information'
     else :
       try:
         policy = json.loads(open(self.policy_path).read())
       except IOError:
         print 'Failed to load policy from %s' % self.policy_path
     return policy
 
-  def HandleRequest(self, path, headers, request):
-    """Handles a request.
-
-    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, machine_id, type):
     """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 = []
@@ skipping 47 matching lines @@
     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_tokens.keys():
       del self._registered_tokens[dmtoken]
+
+
+class PolicyServerRunner(testserver_base.TestServerRunner):
+
+  def __init__(self):
+    super(PolicyServerRunner, self).__init__()
+
+  def create_server(self, server_data):
+    config_file = (
+        self.options.config_file or
+        os.path.join(self.options.data_dir or '', 'device_management'))
+    server = PolicyTestServer((self.options.host, self.options.port),
+                              config_file, self.options.policy_keys)
+    server_data['port'] = server.server_port
+    return server
+
+  def add_options(self):
+    testserver_base.TestServerRunner.add_options(self)
+    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 multipe 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 '
+                                  'generate a random key if none is specified '
+                                  'on the command line.')
+    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')
+
+  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())