tools/telemetry/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 json
  32 import logging
  33 import os
  34
  35 COMMON_URL_PREFIX = '/web-page-replay-'
  36 COMMAND_URL_PREFIX = COMMON_URL_PREFIX + 'command-'
  37 GENERATOR_URL_PREFIX = COMMON_URL_PREFIX + 'generate-'
  38 POST_IMAGE_URL_PREFIX = COMMON_URL_PREFIX + 'post-image-'
  39 IMAGE_DATA_PREFIX = 'data:image/png;base64,'
  40
  41
  42 def SimpleResponse(status):
  43   """Return a ArchivedHttpResponse with |status| code and a simple text body."""
  44   return httparchive.create_response(status)
  45
  46
  47 def JsonResponse(data):
  48   """Return a ArchivedHttpResponse with |data| encoded as json in the body."""
  49   status = 200
  50   reason = 'OK'
  51   headers = [('content-type', 'application/json')]
  52   body = json.dumps(data)
  53   return httparchive.create_response(status, reason, headers, body)
  54
  55
  56 class CustomHandlers(object):
  57
  58   def __init__(self, options, http_archive):
  59     """Initialize CustomHandlers.
  60
  61     Args:
  62       options: original options passed to the server.
  63       http_archive: reference to the HttpArchive object.
  64     """
  65     self.server_manager = None
  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.full_path.startswith(prefix):
  94         return handler(request, request.full_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     Args:
 101       request: an ArchivedHttpRequest instance
 102       url_suffix: string that is after the handler prefix (e.g. 304)
 103     Returns:
 104       On a match, an ArchivedHttpResponse.
 105       Otherwise, None.
 106     """
 107     del request
 108     try:
 109       response_code = int(url_suffix)
 110       return SimpleResponse(response_code)
 111     except ValueError:
 112       return None
 113
 114   def handle_possible_post_image(self, request, url_suffix):
 115     """If sent, saves embedded image to local directory.
 116
 117     Expects a special url containing the filename. If sent, saves the base64
 118     encoded request body as a PNG image locally. This feature is enabled by
 119     passing in screenshot_dir to the initializer for this class.
 120
 121     Args:
 122       request: an ArchivedHttpRequest instance
 123       url_suffix: string that is after the handler prefix (e.g. 'foo.png')
 124     Returns:
 125       On a match, an ArchivedHttpResponse.
 126       Otherwise, None.
 127     """
 128     basename = url_suffix
 129     if not basename:
 130       return None
 131
 132     data = request.request_body
 133     if not data.startswith(IMAGE_DATA_PREFIX):
 134       logging.error('Unexpected image format for: %s', basename)
 135       return SimpleResponse(400)
 136
 137     data = data[len(IMAGE_DATA_PREFIX):]
 138     png = base64.b64decode(data)
 139     filename = os.path.join(self.screenshot_dir,
 140                             '%s-%s.png' % (request.host, basename))
 141     if not os.access(self.screenshot_dir, os.W_OK):
 142       logging.error('Unable to write to: %s', filename)
 143       return SimpleResponse(400)
 144
 145     with file(filename, 'w') as f:
 146       f.write(png)
 147     return SimpleResponse(200)
 148
 149   def add_server_manager_handler(self, server_manager):
 150     """Add the ability to change the server mode (e.g. to record mode).
 151     Args:
 152       server_manager: a servermanager.ServerManager instance.
 153     """
 154     self.server_manager = server_manager
 155     self.handlers.append(
 156         (COMMAND_URL_PREFIX, self.handle_server_manager_command))
 157
 158   def handle_server_manager_command(self, request, url_suffix):
 159     """Parse special URLs for the embedded server manager command.
 160
 161     Clients like webpagetest.org can use URLs of this form to change
 162     the replay server from record mode to replay mode.
 163
 164     This handler is not in the default list of handlers. Call
 165     add_server_manager_handler to add it.
 166
 167     In the future, this could be expanded to save or serve archive files.
 168
 169     Args:
 170       request: an ArchivedHttpRequest instance
 171       url_suffix: string that is after the handler prefix (e.g. 'record')
 172     Returns:
 173       On a match, an ArchivedHttpResponse.
 174       Otherwise, None.
 175     """
 176     command = url_suffix
 177     if command == 'record':
 178       self.server_manager.SetRecordMode()
 179       return SimpleResponse(200)
 180     elif command == 'replay':
 181       self.server_manager.SetReplayMode()
 182       return SimpleResponse(200)
 183     elif command == 'status':
 184       status = {}
 185       is_record_mode = self.server_manager.IsRecordMode()
 186       status['is_record_mode'] = is_record_mode
 187       status['options'] = json.loads(str(self.options))
 188       archive_stats = self.http_archive.stats()
 189       if archive_stats:
 190         status['archive_stats'] = json.loads(archive_stats)
 191       return JsonResponse(status)
 192     elif command == 'exit':
 193       self.server_manager.should_exit = True
 194       return SimpleResponse(200)
 195     elif command == 'log':
 196       logging.info('log command: %s', str(request.request_body)[:1000000])
 197       return SimpleResponse(200)
 198     return None