Check in the thirdparty libs needed for webkitpy.

Tools/Scripts/webkitpy/thirdparty/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 httplib
+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, screenshot_dir=None):
+    """Initialize CustomHandlers.
+
+    Args:
+      screenshot_dir: a path to which screenshots are saved.
+    """
+    self.handlers = [
+        (GENERATOR_URL_PREFIX, self.get_generator_url_response_code)]
+    if screenshot_dir:
+      if not os.path.exists(screenshot_dir):
+        try:
+          os.makedirs(screenshot_dir)
+        except IOError:
+          logging.error('Unable to create screenshot dir: %s', screenshot_dir)
+          screenshot_dir = None
+      if screenshot_dir:
+        self.screenshot_dir = 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.path.startswith(prefix):
+        return handler(request, request.path[len(prefix):])
+    return None
+
+  def get_generator_url_response_code(self, request, url_suffix):
+    """Parse special generator URLs for the embedded response code.
+
+    Clients like perftracker can use URLs of this form to request
+    a response with a particular 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.
+    """
+    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':
+      is_record_mode = self.server_manager.IsRecordMode()
+      return JsonResponse({'is_record_mode': is_record_mode})
+    return None