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 199
