Add -r option to mitmdump and mitmproxy.

This option reads a set of flows from a file. I've also regularized the
mitmdump and mitmproxy command-line signatures by removing mitmproxy's old way
of specifying flow loads through naked arguments.

4 files changed, 26 insertions, 5 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 9e065e4d..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..."
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 3e645167..ce9c2b50 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",
@@ -84,6 +85,15 @@ 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),
@@ -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 b53794c9..14f06241 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -548,13 +548,11 @@ class FlowMaster(controller.Master):
         """
         for i in fr.stream():
             if i.request:
-                f = self.state.add_request(i.request)
-                self.process_new_request(f)
+                self.handle_request(i.request)
             if i.response:
-                f = self.state.add_response(i.response)
-                self.process_new_response(f)
+                self.handle_response(i.response)
             if i.error:
-                f = self.state.add_error(i.error)
+                self.handle_error(i.error)
 
     def process_new_request(self, f):
         if self.stickycookie_state: