Simple record & playback functionality

1 files changed, 273 insertions, 0 deletions

diff --git a/libmproxy/recorder.py b/libmproxy/recorder.py
new file mode 100644
index 00000000..51c8a6e0
--- /dev/null
+++ b/libmproxy/recorder.py
@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+
+# Copyright (C) 2010  Henrik Nordstrom <henrik@henriknordstrom.net>
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# HENRIK NORDSTROM BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
+# OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+# Alternatively you may use this file under a GPLv3 license as follows:
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+# 
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import sys
+import time
+import hashlib
+import utils
+import proxy
+import collections
+import itertools
+import string
+import Cookie
+import filt
+import re
+import cStringIO
+
+def constant_factory(value):
+    return itertools.repeat(value).next
+
+class PatternRule:
+    """
+        Request pattern rule
+        :_ivar  _match          filt pattern rule
+        :_ivar  _search         Regex pattern to search for
+        :_ivar  _replace        Replacement string
+    """
+    def __init__(self, pattern, search, replace):
+        self.match = filt.parse(pattern)
+        self.search = re.compile(search)
+        self.replace = replace
+    def execute(self, request, text):
+        if self.match and not self.match(request):
+            return text
+        return re.sub(self.search, self.replace, text)
+
+class RecorderConnection(proxy.ServerConnection):
+    """
+        Simulated ServerConnection connecting to the cache
+    """
+    # Note: This may chane in future. Division between Recorder
+    # and RecorderConnection is not yet finalized
+    def __init__(self, request, fp):
+        self.host = request.host
+        self.port = request.port
+        self.scheme = request.scheme
+        self.close = False
+        self.server = fp
+        self.rfile = fp
+        self.wfile = fp
+
+    def send_request(self, request):
+        self.request = request
+
+    def read_response(self):
+        response = proxy.ServerConnection.read_response(self)
+        response.cached = True
+        return response
+
+class Recorder:
+    """
+        A simple record/playback cache
+    """
+    def __init__(self, options):
+        self.sequence = collections.defaultdict(int)
+        self.cookies = {}
+        try:
+            for cookie in options.cookies:
+                self.cookies[cookie] = True
+        except AttributeError: pass
+        self.verbosity = options.verbose
+        self.storedir = options.cache
+        self.patterns = []
+        self.indexfp = None
+        self.reset_config()
+
+    def reset_config(self):
+        self.patterns = []
+        self.load_config("default")
+
+    def add_rule(self, match, search, replace):
+        self.patterns.append(PatternRule(match, search, replace))
+
+    def forget_last_rule(self):
+        self.patterns.pop()
+
+    def save_rule(self, match, search, replace, configfile = "default"):
+        fp = self.open(configfile + ".cfg", "a")
+        print >> fp, "Condition: " + match
+        print >> fp, "Search: " + search
+        print >> fp, "Replace: " + replace
+        fp.close()
+
+    def load_config(self, name):
+        """
+            Load configuration settings from name
+        """
+        try:
+            file = name + ".cfg"
+            if self.verbosity > 2:
+                print >> sys.stderr, "config: " + file
+            fp = self.open(file, "r")
+        except IOError:
+            return False
+        for line in fp:
+            directive, value = line.split(" ", 1)
+            value = value.strip("\r\n")
+            if directive == "Cookie:":
+                self.cookies[value] = True
+            if directive == "Condition:":
+                match = value
+            if directive == "Search:":
+                search = value
+            if directive == "Replace:":
+                self.add_rule(match, search, value)
+        fp.close()
+        return True
+
+    def filter_request(self, request):
+        """
+            Filter forwarded requests to enable better recording
+        """
+        request = request.copy()
+        headers = request.headers
+        utils.try_del(headers, 'if-modified-since')
+        utils.try_del(headers, 'if-none-match')
+        return request
+
+    def normalize_request(self, request):
+        """
+            Filter request to simplify storage matching
+        """
+        request.close = False
+        req_text = request.assemble_proxy()
+        orig_req_text = req_text
+        for pattern in self.patterns:
+            req_text = pattern.execute(request, req_text)
+        if req_text == orig_req_text:
+            return request
+        fp = cStringIO.StringIO(req_text)
+        request_line = fp.readline()
+        method, scheme, host, port, path, httpminor = proxy.parse_request_line(request_line)
+        headers = utils.Headers()
+        headers.read(fp)
+        if request.content is None:
+            content = None
+        else:
+            content = fp.read()
+        return proxy.Request(request.client_conn, host, port, scheme, method, path, headers, content)
+
+    def open(self, path, mode):
+        return open(self.storedir + "/" + path, mode)
+
+    def pathn(self, request):
+        """
+            Create cache file name and sequence number
+        """
+        request = self.normalize_request(request)
+        request = self.filter_request(request)
+        headers = request.headers
+        urlkey = (request.host + request.path)[:80].translate(string.maketrans(":/?","__."))
+        id = ""
+        if headers.has_key("cookie"):
+            cookies = Cookie.SimpleCookie("; ".join(headers["cookie"]))
+            del headers["cookie"]
+            for key, morsel in cookies.iteritems():
+                if self.cookies.has_key(key):
+                    id = id + key + "=" + morsel.value + "\n"
+        if self.verbosity > 1:
+            print >> sys.stderr, "ID: " + id
+        m = hashlib.sha224(id)
+        req_text = request.assemble_proxy()
+        if self.verbosity > 2:
+            print >> sys.stderr, req_text
+        m.update(req_text)
+        path = urlkey+"."+m.hexdigest()
+        n = str(self.sequence[path])
+        if self.verbosity > 1:
+            print >> sys.stderr, "PATH: " + path + "." + n
+        return path, n
+
+    def filter_response(self, response):
+        if response.headers.has_key('set-cookie'):
+            for header in response.headers['set-cookie']:
+                key = header.split('=',1)[0]
+                self.cookies[key] = True
+        return response
+
+    def save_response(self, response):
+        """
+            Save response for later playback
+        """
+
+        if self.indexfp is None:
+            self.indexfp = self.open("index.txt", "a")
+            try:
+                cfg = self.open("default.cfg", "r")
+            except:
+                cfg = self.open("default.cfg", "w")
+                for cookie in iter(self.cookies):
+                    print >> cfg, "Cookie: " + cookie
+            cfg.close()
+        request = response.request
+        req_text = request.assemble_proxy()
+        resp_text = response.assemble()
+        path, n = self.pathn(request)
+        self.sequence[path] += 1
+
+        f = self.open(path+"."+n+".req", 'w')
+        f.write(req_text)
+        f.close()
+        f = self.open(path+"."+n+".resp", 'w')
+        f.write(resp_text)
+        f.close()
+
+        print >> self.indexfp , time.time(), request.method, request.path
+        if request.headers.has_key('referer'):
+            print >> self.indexfp, 'referer:', ','.join(request.headers['referer'])
+        if len(self.cookies) > 0:
+            print >> self.indexfp, 'cookies:', ','.join(self.cookies)
+        print >> self.indexfp , path
+        print >> self.indexfp , ""
+
+
+    def get_response(self, request):
+        """
+            Retrieve previously saved response saved by save_response
+        """
+        path, n = self.pathn(request)
+        try:
+            fp = self.open(path+"."+n+".resp", 'r')
+            self.sequence[path]+=1
+        except IOError:
+            fp = self.open(path+".resp", 'r')
+        server = RecorderConnection(request, fp)
+        fp = None       # Handed over to RecorderConnection
+        server.send_request(request)
+        response = server.read_response()
+        server.terminate()
+        return response