device_management.py revision 201ade2fbba22bfb27ae029f4d23fca6ded109a0 
1#!/usr/bin/python2.5
2# Copyright (c) 2010 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""A bare-bones test server for testing cloud policy support.
7
8This implements a simple cloud policy test server that can be used to test
9chrome's device management service client. The policy information is read from
10from files in a directory. The files should contain policy definitions in JSON
11format, using the top-level dictionary as a key/value store. The format is
12identical to what the Linux implementation reads from /etc. Here is an example:
13
14{
15  "HomepageLocation" : "http://www.chromium.org"
16}
17
18"""
19
20import cgi
21import logging
22import random
23import re
24import sys
25
26# The name and availability of the json module varies in python versions.
27try:
28  import simplejson as json
29except ImportError:
30  try:
31    import json
32  except ImportError:
33    json = None
34
35import device_management_backend_pb2 as dm
36
37class RequestHandler(object):
38  """Decodes and handles device management requests from clients.
39
40  The handler implements all the request parsing and protobuf message decoding
41  and encoding. It calls back into the server to lookup, register, and
42  unregister clients.
43  """
44
45  def __init__(self, server, path, headers, request):
46    """Initialize the handler.
47
48    Args:
49      server: The TestServer object to use for (un)registering clients.
50      path: A string containing the request path and query parameters.
51      headers: A rfc822.Message-like object containing HTTP headers.
52      request: The request data received from the client as a string.
53    """
54    self._server = server
55    self._path = path
56    self._headers = headers
57    self._request = request
58    self._params = None
59
60  def GetUniqueParam(self, name):
61    """Extracts a unique query parameter from the request.
62
63    Args:
64      name: Names the parameter to fetch.
65    Returns:
66      The parameter value or None if the parameter doesn't exist or is not
67      unique.
68    """
69    if not self._params:
70      self._params = cgi.parse_qs(self._path[self._path.find('?')+1:])
71
72    param_list = self._params.get(name, [])
73    if len(param_list) == 1:
74      return param_list[0]
75    return None;
76
77  def HandleRequest(self):
78    """Handles a request.
79
80    Parses the data supplied at construction time and returns a pair indicating
81    http status code and response data to be sent back to the client.
82
83    Returns:
84      A tuple of HTTP status code and response data to send to the client.
85    """
86    rmsg = dm.DeviceManagementRequest()
87    rmsg.ParseFromString(self._request)
88
89    self.DumpMessage('Request', rmsg)
90
91    request_type = self.GetUniqueParam('request')
92    if request_type == 'register':
93      return self.ProcessRegister(rmsg.register_request)
94    elif request_type == 'unregister':
95      return self.ProcessUnregister(rmsg.unregister_request)
96    elif request_type == 'policy':
97      return self.ProcessPolicy(rmsg.policy_request)
98    else:
99      return (400, 'Invalid request parameter')
100
101  def ProcessRegister(self, msg):
102    """Handles a register request.
103
104    Checks the query for authorization and device identifier, registers the
105    device with the server and constructs a response.
106
107    Args:
108      msg: The DeviceRegisterRequest message received from the client.
109
110    Returns:
111      A tuple of HTTP status code and response data to send to the client.
112    """
113    # Check the auth token and device ID.
114    match = re.match('GoogleLogin auth=(\\w+)',
115                     self._headers.getheader('Authorization', ''))
116    if not match:
117      return (403, 'No authorization')
118    auth_token = match.group(1)
119
120    device_id = self.GetUniqueParam('deviceid')
121    if not device_id:
122      return (400, 'Missing device identifier')
123
124    # Register the device and create a token.
125    dmtoken = self._server.RegisterDevice(device_id)
126
127    # Send back the reply.
128    response = dm.DeviceManagementResponse()
129    response.error = dm.DeviceManagementResponse.SUCCESS
130    response.register_response.device_management_token = dmtoken
131
132    self.DumpMessage('Response', response)
133
134    return (200, response.SerializeToString())
135
136  def ProcessUnregister(self, msg):
137    """Handles a register request.
138
139    Checks for authorization, unregisters the device and constructs the
140    response.
141
142    Args:
143      msg: The DeviceUnregisterRequest message received from the client.
144
145    Returns:
146      A tuple of HTTP status code and response data to send to the client.
147    """
148    # Check the management token.
149    token, response = self.CheckToken();
150    if not token:
151      return response
152
153    # Unregister the device.
154    self._server.UnregisterDevice(token);
155
156    # Prepare and send the response.
157    response = dm.DeviceManagementResponse()
158    response.error = dm.DeviceManagementResponse.SUCCESS
159    response.unregister_response.CopyFrom(dm.DeviceUnregisterResponse())
160
161    self.DumpMessage('Response', response)
162
163    return (200, response.SerializeToString())
164
165  def ProcessPolicy(self, msg):
166    """Handles a policy request.
167
168    Checks for authorization, encodes the policy into protobuf representation
169    and constructs the repsonse.
170
171    Args:
172      msg: The DevicePolicyRequest message received from the client.
173
174    Returns:
175      A tuple of HTTP status code and response data to send to the client.
176    """
177    # Check the management token.
178    token, response = self.CheckToken()
179    if not token:
180      return response
181
182    # Stuff the policy dictionary into a response message and send it back.
183    response = dm.DeviceManagementResponse()
184    response.error = dm.DeviceManagementResponse.SUCCESS
185    response.policy_response.CopyFrom(dm.DevicePolicyResponse())
186
187    # Respond only if the client requested policy for the cros/device scope,
188    # since that's where chrome policy is supposed to live in.
189    if msg.policy_scope == 'chromeos/device':
190      setting = response.policy_response.setting.add()
191      setting.policy_key = 'chrome-policy'
192      policy_value = dm.GenericSetting()
193      for (key, value) in self._server.policy.iteritems():
194        entry = policy_value.named_value.add()
195        entry.name = key
196        entry_value = dm.GenericValue()
197        if isinstance(value, bool):
198          entry_value.value_type = dm.GenericValue.VALUE_TYPE_BOOL
199          entry_value.bool_value = value
200        elif isinstance(value, int):
201          entry_value.value_type = dm.GenericValue.VALUE_TYPE_INT64
202          entry_value.int64_value = value
203        elif isinstance(value, str) or isinstance(value, unicode):
204          entry_value.value_type = dm.GenericValue.VALUE_TYPE_STRING
205          entry_value.string_value = value
206        elif isinstance(value, list):
207          entry_value.value_type = dm.GenericValue.VALUE_TYPE_STRING_ARRAY
208          for list_entry in value:
209            entry_value.string_array.append(str(list_entry))
210        entry.value.CopyFrom(entry_value)
211      setting.policy_value.CopyFrom(policy_value)
212
213    self.DumpMessage('Response', response)
214
215    return (200, response.SerializeToString())
216
217  def CheckToken(self):
218    """Helper for checking whether the client supplied a valid DM token.
219
220    Extracts the token from the request and passed to the server in order to
221    look up the client. Returns a pair of token and error response. If the token
222    is None, the error response is a pair of status code and error message.
223
224    Returns:
225      A pair of DM token and error response. If the token is None, the message
226      will contain the error response to send back.
227    """
228    error = None
229    dmtoken = None
230    request_device_id = self.GetUniqueParam('deviceid')
231    match = re.match('GoogleDMToken token=(\\w+)',
232                     self._headers.getheader('Authorization', ''))
233    if match:
234      dmtoken = match.group(1)
235    if not dmtoken:
236      error = dm.DeviceManagementResponse.DEVICE_MANAGEMENT_TOKEN_INVALID
237    elif (not request_device_id or
238          not self._server.LookupDevice(dmtoken) == request_device_id):
239      error = dm.DeviceManagementResponse.DEVICE_NOT_FOUND
240    else:
241      return (dmtoken, None)
242
243    response = dm.DeviceManagementResponse()
244    response.error = error
245
246    self.DumpMessage('Response', response)
247
248    return (None, (200, response.SerializeToString()))
249
250  def DumpMessage(self, label, msg):
251    """Helper for logging an ASCII dump of a protobuf message."""
252    logging.debug('%s\n%s' % (label, str(msg)))
253
254class TestServer(object):
255  """Handles requests and keeps global service state."""
256
257  def __init__(self, policy_path):
258    """Initializes the server.
259
260    Args:
261      policy_path: Names the file to read JSON-formatted policy from.
262    """
263    self._registered_devices = {}
264    self.policy = {}
265    if json is None:
266      print 'No JSON module, cannot parse policy information'
267    else :
268      try:
269        self.policy = json.loads(open(policy_path).read())
270      except IOError:
271        print 'Failed to load policy from %s' % policy_path
272
273  def HandleRequest(self, path, headers, request):
274    """Handles a request.
275
276    Args:
277      path: The request path and query parameters received from the client.
278      headers: A rfc822.Message-like object containing HTTP headers.
279      request: The request data received from the client as a string.
280    Returns:
281      A pair of HTTP status code and response data to send to the client.
282    """
283    handler = RequestHandler(self, path, headers, request)
284    return handler.HandleRequest()
285
286  def RegisterDevice(self, device_id):
287    """Registers a device and generate a DM token for it.
288
289    Args:
290      device_id: The device identifier provided by the client.
291
292    Returns:
293      The newly generated device token for the device.
294    """
295    dmtoken_chars = []
296    while len(dmtoken_chars) < 32:
297      dmtoken_chars.append(random.choice('0123456789abcdef'))
298    dmtoken= ''.join(dmtoken_chars)
299    self._registered_devices[dmtoken] = device_id
300    return dmtoken
301
302  def LookupDevice(self, dmtoken):
303    """Looks up a device by DMToken.
304
305    Args:
306      dmtoken: The device management token provided by the client.
307
308    Returns:
309      The corresponding device identifier or None if not found.
310    """
311    return self._registered_devices.get(dmtoken, None)
312
313  def UnregisterDevice(self, dmtoken):
314    """Unregisters a device identified by the given DM token.
315
316    Args:
317      dmtoken: The device management token provided by the client.
318    """
319    if dmtoken in self._registered_devices:
320      del self._registered_devices[dmtoken]
321