• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
29import base64
30import httparchive
31import json
32import logging
33import os
34
35COMMON_URL_PREFIX = '/web-page-replay-'
36COMMAND_URL_PREFIX = COMMON_URL_PREFIX + 'command-'
37GENERATOR_URL_PREFIX = COMMON_URL_PREFIX + 'generate-'
38POST_IMAGE_URL_PREFIX = COMMON_URL_PREFIX + 'post-image-'
39IMAGE_DATA_PREFIX = 'data:image/png;base64,'
40
41
42def SimpleResponse(status):
43  """Return a ArchivedHttpResponse with |status| code and a simple text body."""
44  return httparchive.create_response(status)
45
46
47def 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
56class 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
199