src/third_party/webpagereplay/customhandlers.py

   1 #!/usr/bin/env python
   2 # Copyright 2010 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """Handle special HTTP requests.
  17
  18 /web-page-replay-generate-[RESPONSE_CODE]
  19   - Return the given RESPONSE_CODE.
  20 /web-page-replay-post-image-[FILENAME]
  21   - Save the posted image to local disk.
  22 /web-page-replay-command-[record|replay|status]
  23   - Optional. Enable by calling custom_handlers.add_server_manager_handler(...).
  24   - Change the server mode to either record or replay.
  25     + When switching to record, the http_archive is cleared.
  26     + When switching to replay, the http_archive is maintained.
  27 """
  28
  29 import base64
  30 import httparchive
  31 import httplib
  32 import json
  33 import logging
  34 import os
  35
  36 COMMON_URL_PREFIX = '/web-page-replay-'
  37 COMMAND_URL_PREFIX = COMMON_URL_PREFIX + 'command-'
  38 GENERATOR_URL_PREFIX = COMMON_URL_PREFIX + 'generate-'
  39 POST_IMAGE_URL_PREFIX = COMMON_URL_PREFIX + 'post-image-'
  40 IMAGE_DATA_PREFIX = 'data:image/png;base64,'
  41
  42
  43 def SimpleResponse(status):
  44   """Return a ArchivedHttpResponse with |status| code and a simple text body."""
  45   return httparchive.create_response(status)
  46
  47
  48 def JsonResponse(data):
  49   """Return a ArchivedHttpResponse with |data| encoded as json in the body."""
  50   status = 200
  51   reason = 'OK'
  52   headers = [('content-type', 'application/json')]
  53   body = json.dumps(data)
  54   return httparchive.create_response(status, reason, headers, body)
  55
  56
  57 class CustomHandlers(object):
  58
  59   def __init__(self, options, http_archive):
  60     """Initialize CustomHandlers.
  61
  62     Args:
  63       options: original options passed to the server.
  64       http_archive: reference to the HttpArchive object.
  65     """
  66     self.options = options
  67     self.http_archive = http_archive
  68     self.handlers = [
  69         (GENERATOR_URL_PREFIX, self.get_generator_url_response_code)]
  70     # screenshot_dir is a path to which screenshots are saved.
  71     if options.screenshot_dir:
  72       if not os.path.exists(options.screenshot_dir):
  73         try:
  74           os.makedirs(options.screenshot_dir)
  75         except IOError:
  76           logging.error('Unable to create screenshot dir: %s',
  77                          options.screenshot_dir)
  78           options.screenshot_dir = None
  79       if options.screenshot_dir:
  80         self.screenshot_dir = options.screenshot_dir
  81         self.handlers.append(
  82             (POST_IMAGE_URL_PREFIX, self.handle_possible_post_image))
  83
  84   def handle(self, request):
  85     """Dispatches requests to matching handlers.
  86
  87     Args:
  88       request: an http request
  89     Returns:
  90       ArchivedHttpResponse or None.
  91     """
  92     for prefix, handler in self.handlers:
  93       if request.path.startswith(prefix):
  94         return handler(request, request.path[len(prefix):])
  95     return None
  96
  97   def get_generator_url_response_code(self, request, url_suffix):
  98     """Parse special generator URLs for the embedded response code.
  99
 100     Clients like perftracker can use URLs of this form to request
 101     a response with a particular response code.
 102
 103     Args:
 104       request: an ArchivedHttpRequest instance
 105       url_suffix: string that is after the handler prefix (e.g. 304)
 106     Returns:
 107       On a match, an ArchivedHttpResponse.
 108       Otherwise, None.
 109     """
 110     try:
 111       response_code = int(url_suffix)
 112       return SimpleResponse(response_code)
 113     except ValueError:
 114       return None
 115
 116   def handle_possible_post_image(self, request, url_suffix):
 117     """If sent, saves embedded image to local directory.
 118
 119     Expects a special url containing the filename. If sent, saves the base64
 120     encoded request body as a PNG image locally. This feature is enabled by
 121     passing in screenshot_dir to the initializer for this class.
 122
 123     Args:
 124       request: an ArchivedHttpRequest instance
 125       url_suffix: string that is after the handler prefix (e.g. 'foo.png')
 126     Returns:
 127       On a match, an ArchivedHttpResponse.
 128       Otherwise, None.
 129     """
 130     basename = url_suffix
 131     if not basename:
 132       return None
 133
 134     data = request.request_body
 135     if not data.startswith(IMAGE_DATA_PREFIX):
 136       logging.error('Unexpected image format for: %s', basename)
 137       return SimpleResponse(400)
 138
 139     data = data[len(IMAGE_DATA_PREFIX):]
 140     png = base64.b64decode(data)
 141     filename = os.path.join(self.screenshot_dir,
 142                             '%s-%s.png' % (request.host, basename))
 143     if not os.access(self.screenshot_dir, os.W_OK):
 144       logging.error('Unable to write to: %s', filename)
 145       return SimpleResponse(400)
 146
 147     with file(filename, 'w') as f:
 148       f.write(png)
 149     return SimpleResponse(200)
 150
 151   def add_server_manager_handler(self, server_manager):
 152     """Add the ability to change the server mode (e.g. to record mode).
 153     Args:
 154       server_manager: a servermanager.ServerManager instance.
 155     """
 156     self.server_manager = server_manager
 157     self.handlers.append(
 158         (COMMAND_URL_PREFIX, self.handle_server_manager_command))
 159
 160   def handle_server_manager_command(self, request, url_suffix):
 161     """Parse special URLs for the embedded server manager command.
 162
 163     Clients like webpagetest.org can use URLs of this form to change
 164     the replay server from record mode to replay mode.
 165
 166     This handler is not in the default list of handlers. Call
 167     add_server_manager_handler to add it.
 168
 169     In the future, this could be expanded to save or serve archive files.
 170
 171     Args:
 172       request: an ArchivedHttpRequest instance
 173       url_suffix: string that is after the handler prefix (e.g. 'record')
 174     Returns:
 175       On a match, an ArchivedHttpResponse.
 176       Otherwise, None.
 177     """
 178     command = url_suffix
 179     if command == 'record':
 180       self.server_manager.SetRecordMode()
 181       return SimpleResponse(200)
 182     elif command == 'replay':
 183       self.server_manager.SetReplayMode()
 184       return SimpleResponse(200)
 185     elif command == 'status':
 186       status = {}
 187       is_record_mode = self.server_manager.IsRecordMode()
 188       status['is_record_mode'] = is_record_mode
 189       status['options'] = json.loads(str(self.options))
 190       archive_stats = self.http_archive.stats()
 191       if archive_stats:
 192         status['archive_stats'] = json.loads(archive_stats)
 193       return JsonResponse(status)
 194     elif command == 'exit':
 195       self.server_manager.should_exit = True
 196       return SimpleResponse(200)
 197     elif command == 'log':
 198       logging.info('log command: %s', str(request.request_body)[:1000000])
 199       return SimpleResponse(200)
 200     return None