From e285b17e3ff2e48455a5bc45b9a401b5b78b65cf Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 15 May 2011 11:54:12 +1200
Subject: [PATCH] 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.
---
 libmproxy/cmdline.py |  6 ++++++
 libmproxy/console.py |  5 +++++
 libmproxy/dump.py    | 12 ++++++++++++
 libmproxy/flow.py    |  8 +++-----
 mitmdump             |  2 +-
 mitmproxy            |  9 +--------
 6 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index 222defebf..ce68baed9 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,
@@ -57,6 +58,11 @@ def common_options(parser):
         action="store_true", dest="quiet",
         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,
diff --git a/libmproxy/console.py b/libmproxy/console.py
index 9e065e4d8..aa7639301 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 3e645167f..ce9c2b50c 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 b53794c9b..14f06241a 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:
diff --git a/mitmdump b/mitmdump
index b54d051f6..bb4ec5017 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 bcb079a30..942646b22 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()