Merge branch 'master' of github.com:cortesi/mitmproxy

8 files changed, 76 insertions, 18 deletions

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index 222defeb..ce68baed 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -21,6 +21,7 @@ def get_common_options(options):
         no_server = options.no_server,
         refresh_server_playback = not options.norefresh,
         rheaders = options.rheaders,
+        rfile = options.rfile,
         request_script = options.request_script,
         response_script = options.response_script,
         server_replay = options.server_replay,
@@ -58,6 +59,11 @@ def common_options(parser):
         help="Quiet."
     )
     parser.add_option(
+        "-r",
+        action="store", dest="rfile", default=None,
+        help="Read flows from file."
+    )
+    parser.add_option(
         "--anticache",
         action="store_true", dest="anticache", default=False,
         help="Strip out request headers that might cause the server to return 304-not-modified."
diff --git a/libmproxy/console.py b/libmproxy/console.py
index 63dbf4f0..aa763930 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -782,6 +782,7 @@ class Options(object):
         "refresh_server_playback",
         "request_script",
         "response_script",
+        "rfile",
         "rheaders",
         "server_replay",
         "stickycookie",
@@ -818,6 +819,7 @@ class ConsoleMaster(flow.FlowMaster):
     def __init__(self, server, options):
         flow.FlowMaster.__init__(self, server, ConsoleState())
         self.looptime = 0
+        self.options = options
 
         self.conn_list_view = None
         self.set_palette()
@@ -1031,6 +1033,9 @@ class ConsoleMaster(flow.FlowMaster):
             slave = controller.Slave(self.masterq, self.server)
             slave.start()
 
+        if self.options.rfile:
+            self.load_flows(self.options.rfile)
+
         self.ui.run_wrapper(self.loop)
         # If True, quit just pops out to connection list view.
         print >> sys.stderr, "Shutting down..."
@@ -1127,11 +1132,10 @@ class ConsoleMaster(flow.FlowMaster):
         try:
             f = file(path, "r")
             fr = flow.FlowReader(f)
-            data = list(fr.stream())
-            f.close()
         except IOError, v:
             return v.strerror
-        self.state.load_flows(data)
+        flow.FlowMaster.load_flows(self, fr)
+        f.close()
         if self.conn_list_view:
             self.sync_list_view()
             self.focus_current()
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 3e645167..b4621716 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -14,6 +14,7 @@ class Options(object):
         "refresh_server_playback",
         "request_script",
         "response_script",
+        "rfile",
         "rheaders",
         "server_replay",
         "stickycookie",
@@ -52,6 +53,8 @@ class DumpMaster(flow.FlowMaster):
         flow.FlowMaster.__init__(self, server, flow.State())
         self.outfile = outfile
         self.o = options
+        self.anticache = options.anticache
+        self.refresh_server_playback = options.refresh_server_playback
 
         if filtstr:
             self.filt = filt.parse(filtstr)
@@ -84,14 +87,21 @@ class DumpMaster(flow.FlowMaster):
                 not options.keepserving
             )
 
+        if options.rfile:
+            path = os.path.expanduser(options.rfile)
+            try:
+                f = file(path, "r")
+                freader = flow.FlowReader(f)
+            except IOError, v:
+                raise DumpError(v.strerror)
+            self.load_flows(freader)
+
         if options.client_replay:
             self.start_client_playback(
                 self._readflow(options.client_replay),
                 not options.keepserving
             )
 
-        self.anticache = options.anticache
-        self.refresh_server_playback = options.refresh_server_playback
 
     def _readflow(self, path):
         path = os.path.expanduser(path)
@@ -189,6 +199,8 @@ class DumpMaster(flow.FlowMaster):
 
 # begin nocover
     def run(self):
+        if self.o.rfile and not self.o.keepserving:
+            return
         try:
             return flow.FlowMaster.run(self)
         except BaseException, v:
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 387c49f0..14f06241 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -542,6 +542,18 @@ class FlowMaster(controller.Master):
 
         controller.Master.tick(self, q)
 
+    def load_flows(self, fr):
+        """
+            Load flows from a FlowReader object.
+        """
+        for i in fr.stream():
+            if i.request:
+                self.handle_request(i.request)
+            if i.response:
+                self.handle_response(i.response)
+            if i.error:
+                self.handle_error(i.error)
+
     def process_new_request(self, f):
         if self.stickycookie_state:
             self.stickycookie_state.handle_request(f)
@@ -560,6 +572,12 @@ class FlowMaster(controller.Master):
                 else:
                     f.request.ack()
 
+    def process_new_response(self, f):
+        if self.stickycookie_state:
+            self.stickycookie_state.handle_response(f)
+        if "response" in self.scripts:
+            self._runscript(f, self.scripts["response"])
+
     def replay_request(self, f):
         """
             Returns None if successful, or error message if not.
@@ -605,10 +623,7 @@ class FlowMaster(controller.Master):
             self.client_playback.clear(f)
         if not f:
             r.ack()
-        if self.stickycookie_state:
-            self.stickycookie_state.handle_response(f)
-        if "response" in self.scripts:
-            self._runscript(f, self.scripts["response"])
+        self.process_new_response(f)
         return f
 
 
diff --git a/mitmdump b/mitmdump
index b54d051f..bb4ec501 100755
--- a/mitmdump
+++ b/mitmdump
@@ -30,7 +30,7 @@ if __name__ == '__main__':
     parser.add_option(
         "--keepserving",
         action="store_true", dest="keepserving", default=False,
-        help="Continue serving after playback. We exit by default."
+        help="Continue serving after client playback or file read. We exit by default."
     )
 
     options, args = parser.parse_args()
diff --git a/mitmproxy b/mitmproxy
index bcb079a3..942646b2 100755
--- a/mitmproxy
+++ b/mitmproxy
@@ -23,7 +23,7 @@ from optparse import OptionParser, OptionGroup
 
 if __name__ == '__main__':
     parser = OptionParser(
-                usage = "%prog [options] [flowdump path]",
+                usage = "%prog [options]",
                 version="%%prog %s"%VERSION,
             )
     cmdline.common_options(parser)
@@ -65,13 +65,6 @@ if __name__ == '__main__':
     opts.debug = options.debug
     m = console.ConsoleMaster(server, opts)
 
-    for i in args:
-        try:
-            m.load_flows(i)
-        except flow.FlowReadError, v:
-            print >> sys.stderr, "mitmproxy:", v.args[0]
-            sys.exit(1)
-
     m.run()
 
 
diff --git a/test/test_flow.py b/test/test_flow.py
index a6489462..b7a368d2 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -369,6 +369,19 @@ class uState(libpry.AutoTree):
 
 
 class uSerialize(libpry.AutoTree):
+    def _treader(self):
+        sio = StringIO()
+        w = flow.FlowWriter(sio)
+        for i in range(3):
+            f = tutils.tflow_full()
+            w.add(f)
+        for i in range(3):
+            f = tutils.tflow_err()
+            w.add(f)
+
+        sio.seek(0)
+        return flow.FlowReader(sio)
+
     def test_roundtrip(self):
         sio = StringIO()
         f = tutils.tflow()
@@ -381,6 +394,14 @@ class uSerialize(libpry.AutoTree):
         assert len(l) == 1
         assert l[0] == f
 
+    def test_load_flows(self):
+        r = self._treader()
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        fm.load_flows(r)
+        assert len(s.flow_list) == 6
+
+
     def test_error(self):
         sio = StringIO()
         sio.write("bogus")
diff --git a/test/tutils.py b/test/tutils.py
index fea7f224..2327bdb9 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -32,6 +32,13 @@ def tflow_full():
     return f
 
 
+def tflow_err():
+    r = treq()
+    f = flow.Flow(r)
+    f.error = proxy.Error(r, "error")
+    return f
+
+
 # Yes, the random ports are horrible. During development, sockets are often not
 # properly closed during error conditions, which means you have to wait until
 # you can re-bind to the same port. This is a pain in the ass, so we just pick