From 061cea89da96bc2ee71d9d2f065ae920aeccb311 Mon Sep 17 00:00:00 2001
From: Henrik Nordstrom <henrik@henriknordstrom.net>
Date: Tue, 11 Jan 2011 18:04:15 +0100
Subject: [PATCH] Import cache store control into console controller

---
 libmproxy/console.py  | 10 ++++++++++
 libmproxy/recorder.py |  6 +++++-
 libmproxy/utils.py    | 11 ++++++++++-
 mitmproxy             | 22 ++++++++++++++++++++++
 mitmrecord            | 13 ++-----------
 5 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 1f19671ae..9bf032e68 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -19,6 +19,7 @@ import cStringIO
 import urwid.curses_display
 import urwid
 import controller, utils, filt, proxy, flow
+import recorder
 
 
 class Stop(Exception): pass
@@ -707,6 +708,7 @@ class ConsoleState(flow.State):
         flow.State.__init__(self)
         self.focus = None
         self.beep = None
+        self.store = None
 
         self.view_body_mode = VIEW_BODY_RAW
         self.view_flow_mode = VIEW_FLOW_REQUEST
@@ -727,6 +729,8 @@ class ConsoleState(flow.State):
         return flow.State.add_request(self, req)
 
     def add_response(self, resp):
+        if self.store is not None:
+            self.store.save_response(resp)
         f = flow.State.add_response(self, resp)
         if self.focus is None:
             self.set_focus(0)
@@ -737,6 +741,9 @@ class ConsoleState(flow.State):
         self.set_focus(self.focus)
         return ret
 
+    def start_recording(self, recorder):
+        self.store = recorder
+
     def get_focus(self):
         if not self.view or self.focus is None:
             return None, None
@@ -815,6 +822,9 @@ class ConsoleMaster(controller.Master):
         self.stickycookie = None
         self.stickyhosts = {}
 
+        if options.cache is not None:
+            self.state.start_recording(recorder.Recorder(options))
+
     def spawn_external_viewer(self, data, contenttype):
         if contenttype:
             ext = mimetypes.guess_extension(contenttype) or ""
diff --git a/libmproxy/recorder.py b/libmproxy/recorder.py
index 51c8a6e02..bbbe36646 100644
--- a/libmproxy/recorder.py
+++ b/libmproxy/recorder.py
@@ -101,7 +101,10 @@ class Recorder:
             for cookie in options.cookies:
                 self.cookies[cookie] = True
         except AttributeError: pass
-        self.verbosity = options.verbose
+        try:
+            self.verbosity = options.verbose
+        except AttributeError:
+            self.verbosity = False
         self.storedir = options.cache
         self.patterns = []
         self.indexfp = None
@@ -253,6 +256,7 @@ class Recorder:
             print >> self.indexfp, 'cookies:', ','.join(self.cookies)
         print >> self.indexfp , path
         print >> self.indexfp , ""
+        self.indexfp.flush()
 
 
     def get_response(self, request):
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 42e4d28c6..87fca5ce6 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -12,7 +12,7 @@
 # 
 # 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 re, os, subprocess, datetime, textwrap
+import re, os, subprocess, datetime, textwrap, errno
 
 
 def format_timestamp(s):
@@ -341,3 +341,12 @@ def make_bogus_cert(path):
         stdin=subprocess.PIPE
     )
     
+def mkdir_p(path):
+    try:
+        os.makedirs(path)
+    except OSError as exc:
+        if exc.errno == errno.EEXIST:
+            pass
+        else:
+            raise
+
diff --git a/mitmproxy b/mitmproxy
index 7b55b8489..074a439c6 100755
--- a/mitmproxy
+++ b/mitmproxy
@@ -72,6 +72,23 @@ if __name__ == '__main__':
     )
     parser.add_option_group(group)
 
+    group = OptionGroup(
+                        parser,
+                        "Recorder",
+                        "Options controlling recorder behavior"
+                    )
+    group.add_option(
+        "-w", "--store", action="store",
+        type = "str", dest="cache", default=None,
+        help = "Session store location"
+    )
+    group.add_option(
+        "-C", "--cookies", action="append",
+        type = "str", dest="cookies", default=[],
+        help = "Persistent client cookies already set or generated in client"
+    )
+    parser.add_option_group(group)
+
     options, args = parser.parse_args()
     certpath = os.path.expanduser(options.cert)
     
@@ -82,6 +99,11 @@ if __name__ == '__main__':
     proxy.config = proxy.Config(
         certpath
     )
+    if options.cache is not None:
+        utils.mkdir_p(options.cache)
+        if os.path.exists(options.cache + "/index.txt"):
+            print >> sys.stderr, "ERROR: data already recorded in %s"%options.cache
+            sys.exit(1)
     server = proxy.ProxyServer(options.port, options.addr)
     m = console.ConsoleMaster(server, options)
 
diff --git a/mitmrecord b/mitmrecord
index 0fcb91c2c..1fcfde4a3 100755
--- a/mitmrecord
+++ b/mitmrecord
@@ -18,20 +18,11 @@
 # 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, os.path, os, errno
+import sys, os.path
 from libmproxy import proxy, controller, record, utils
 from libmproxy import VERSION
 from optparse import OptionParser, OptionGroup
 
-def mkdir_p(path):
-    try:
-        os.makedirs(path)
-    except OSError as exc:
-        if exc.errno == errno.EEXIST:
-            pass
-        else:
-            raise
-
 if __name__ == '__main__':
     parser = OptionParser(
                 usage = "%prog [options] output",
@@ -85,7 +76,7 @@ if __name__ == '__main__':
         certpath
     )
     server = proxy.ProxyServer(options.port)
-    mkdir_p(options.cache)
+    utils.mkdir_p(options.cache)
     if os.path.exists(options.cache + "/index.txt"):
         print >> sys.stderr, "ERROR: data already recorded in %s"%options.cache
         sys.exit(1)