Rename third_party/webpagereplay to third_party/web-page-replay

telemetry/third_party/webpagereplay/customhandlers.py

-#!/usr/bin/env python
-# Copyright 2010 Google Inc. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Handle special HTTP requests.
-
-/web-page-replay-generate-[RESPONSE_CODE]
-  - Return the given RESPONSE_CODE.
-/web-page-replay-post-image-[FILENAME]
-  - Save the posted image to local disk.
-/web-page-replay-command-[record|replay|status]
-  - Optional. Enable by calling custom_handlers.add_server_manager_handler(...).
-  - Change the server mode to either record or replay.
-    + When switching to record, the http_archive is cleared.
-    + When switching to replay, the http_archive is maintained.
-"""
-
-import base64
-import httparchive
-import json
-import logging
-import os
-
-COMMON_URL_PREFIX = '/web-page-replay-'
-COMMAND_URL_PREFIX = COMMON_URL_PREFIX + 'command-'
-GENERATOR_URL_PREFIX = COMMON_URL_PREFIX + 'generate-'
-POST_IMAGE_URL_PREFIX = COMMON_URL_PREFIX + 'post-image-'
-IMAGE_DATA_PREFIX = 'data:image/png;base64,'
-
-
-def SimpleResponse(status):
-  """Return a ArchivedHttpResponse with |status| code and a simple text body."""
-  return httparchive.create_response(status)
-
-
-def JsonResponse(data):
-  """Return a ArchivedHttpResponse with |data| encoded as json in the body."""
-  status = 200
-  reason = 'OK'
-  headers = [('content-type', 'application/json')]
-  body = json.dumps(data)
-  return httparchive.create_response(status, reason, headers, body)
-
-
-class CustomHandlers(object):
-
-  def __init__(self, options, http_archive):
-    """Initialize CustomHandlers.
-
-    Args:
-      options: original options passed to the server.
-      http_archive: reference to the HttpArchive object.
-    """
-    self.server_manager = None
-    self.options = options
-    self.http_archive = http_archive
-    self.handlers = [
-        (GENERATOR_URL_PREFIX, self.get_generator_url_response_code)]
-    # screenshot_dir is a path to which screenshots are saved.
-    if options.screenshot_dir:
-      if not os.path.exists(options.screenshot_dir):
-        try:
-          os.makedirs(options.screenshot_dir)
-        except IOError:
-          logging.error('Unable to create screenshot dir: %s',
-                         options.screenshot_dir)
-          options.screenshot_dir = None
-      if options.screenshot_dir:
-        self.screenshot_dir = options.screenshot_dir
-        self.handlers.append(
-            (POST_IMAGE_URL_PREFIX, self.handle_possible_post_image))
-
-  def handle(self, request):
-    """Dispatches requests to matching handlers.
-
-    Args:
-      request: an http request
-    Returns:
-      ArchivedHttpResponse or None.
-    """
-    for prefix, handler in self.handlers:
-      if request.full_path.startswith(prefix):
-        return handler(request, request.full_path[len(prefix):])
-    return None
-
-  def get_generator_url_response_code(self, request, url_suffix):
-    """Parse special generator URLs for the embedded response code.
-
-    Args:
-      request: an ArchivedHttpRequest instance
-      url_suffix: string that is after the handler prefix (e.g. 304)
-    Returns:
-      On a match, an ArchivedHttpResponse.
-      Otherwise, None.
-    """
-    del request
-    try:
-      response_code = int(url_suffix)
-      return SimpleResponse(response_code)
-    except ValueError:
-      return None
-
-  def handle_possible_post_image(self, request, url_suffix):
-    """If sent, saves embedded image to local directory.
-
-    Expects a special url containing the filename. If sent, saves the base64
-    encoded request body as a PNG image locally. This feature is enabled by
-    passing in screenshot_dir to the initializer for this class.
-
-    Args:
-      request: an ArchivedHttpRequest instance
-      url_suffix: string that is after the handler prefix (e.g. 'foo.png')
-    Returns:
-      On a match, an ArchivedHttpResponse.
-      Otherwise, None.
-    """
-    basename = url_suffix
-    if not basename:
-      return None
-
-    data = request.request_body
-    if not data.startswith(IMAGE_DATA_PREFIX):
-      logging.error('Unexpected image format for: %s', basename)
-      return SimpleResponse(400)
-
-    data = data[len(IMAGE_DATA_PREFIX):]
-    png = base64.b64decode(data)
-    filename = os.path.join(self.screenshot_dir,
-                            '%s-%s.png' % (request.host, basename))
-    if not os.access(self.screenshot_dir, os.W_OK):
-      logging.error('Unable to write to: %s', filename)
-      return SimpleResponse(400)
-
-    with file(filename, 'w') as f:
-      f.write(png)
-    return SimpleResponse(200)
-
-  def add_server_manager_handler(self, server_manager):
-    """Add the ability to change the server mode (e.g. to record mode).
-    Args:
-      server_manager: a servermanager.ServerManager instance.
-    """
-    self.server_manager = server_manager
-    self.handlers.append(
-        (COMMAND_URL_PREFIX, self.handle_server_manager_command))
-
-  def handle_server_manager_command(self, request, url_suffix):
-    """Parse special URLs for the embedded server manager command.
-
-    Clients like webpagetest.org can use URLs of this form to change
-    the replay server from record mode to replay mode.
-
-    This handler is not in the default list of handlers. Call
-    add_server_manager_handler to add it.
-
-    In the future, this could be expanded to save or serve archive files.
-
-    Args:
-      request: an ArchivedHttpRequest instance
-      url_suffix: string that is after the handler prefix (e.g. 'record')
-    Returns:
-      On a match, an ArchivedHttpResponse.
-      Otherwise, None.
-    """
-    command = url_suffix
-    if command == 'record':
-      self.server_manager.SetRecordMode()
-      return SimpleResponse(200)
-    elif command == 'replay':
-      self.server_manager.SetReplayMode()
-      return SimpleResponse(200)
-    elif command == 'status':
-      status = {}
-      is_record_mode = self.server_manager.IsRecordMode()
-      status['is_record_mode'] = is_record_mode
-      status['options'] = json.loads(str(self.options))
-      archive_stats = self.http_archive.stats()
-      if archive_stats:
-        status['archive_stats'] = json.loads(archive_stats)
-      return JsonResponse(status)
-    elif command == 'exit':
-      self.server_manager.should_exit = True
-      return SimpleResponse(200)
-    elif command == 'log':
-      logging.info('log command: %s', str(request.request_body)[:1000000])
-      return SimpleResponse(200)
-    return None