From a7e64a1a03b8ebf424569763c84dfb2209a3b96c Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 6 Jul 2012 14:41:10 +1200
Subject: mitmproxy: "W" shortcut key streams flows to file as responses
 arrive.

---
 libmproxy/console/__init__.py | 29 +++++++++++++++++++++++++++--
 libmproxy/console/flowlist.py | 15 +++++++++++----
 2 files changed, 38 insertions(+), 6 deletions(-)

(limited to 'libmproxy/console')

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 8574b93f..5501b935 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -25,8 +25,6 @@ EVENTLOG_SIZE = 500
 class Stop(Exception): pass
 
 
-
-
 class _PathCompleter:
     def __init__(self, _testing=False):
         """
@@ -187,6 +185,9 @@ class StatusBar(common.WWrap):
         if self.master.debug:
             r.append("[lt:%0.3f]"%self.master.looptime)
 
+        if self.master.stream:
+            r.append("[W:%s]"%self.master.stream_path)
+
         return r
 
     def redraw(self):
@@ -353,6 +354,7 @@ class ConsoleMaster(flow.FlowMaster):
     def __init__(self, server, options):
         flow.FlowMaster.__init__(self, server, ConsoleState())
         self.looptime = 0
+        self.stream = None
         self.options = options
 
         for i in options.replacements:
@@ -400,6 +402,25 @@ class ConsoleMaster(flow.FlowMaster):
                 print >> sys.stderr, "Script load error:", err
                 sys.exit(1)
 
+        if options.wfile:
+            err = self.start_stream(options.wfile)
+            if err:
+                print >> sys.stderr, "Script load error:", err
+                sys.exit(1)
+
+    def start_stream(self, path):
+        path = os.path.expanduser(path)
+        try:
+            f = file(path, "ab")
+            self.stream = flow.FlowWriter(f)
+        except IOError, v:
+            return str(v)
+        self.stream_path = path
+
+    def stop_stream(self):
+        self.stream = None
+        self.stream_path = None
+
     def run_script_once(self, path, f):
         if not path:
             return
@@ -905,6 +926,8 @@ class ConsoleMaster(flow.FlowMaster):
 
     def shutdown(self):
         self.state.killall(self)
+        if self.stream:
+            self.stream.fo.close()
         controller.Master.shutdown(self)
 
     def sync_list_view(self):
@@ -969,4 +992,6 @@ class ConsoleMaster(flow.FlowMaster):
         f = flow.FlowMaster.handle_response(self, r)
         if f:
             self.process_flow(f, r)
+        if self.stream:
+            self.stream.add(f)
         return f
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index d80849e1..71214e5e 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -29,8 +29,8 @@ def _mkhelp():
         ("L", "load saved flows"),
         ("r", "replay request"),
         ("V", "revert changes to request"),
-        ("w", "save all flows matching current limit"),
-        ("W", "save this flow"),
+        ("w", "save flows "),
+        ("W", "stream flows to file"),
         ("X", "kill and delete flow, even if it's mid-intercept"),
         ("tab", "tab between eventlog and flow list"),
         ("enter", "view flow"),
@@ -161,8 +161,6 @@ class ConnectionItem(common.WWrap):
                 ),
                 self.save_flows_prompt,
             )
-        elif key == "W":
-            pass
         elif key == "X":
             self.flow.kill(self.master)
         elif key == "enter":
@@ -227,5 +225,14 @@ class FlowListBox(urwid.ListBox):
                 self.master.state.last_saveload,
                 self.master.load_flows_callback
             )
+        elif key == "W":
+            if self.master.stream:
+                self.master.stop_stream()
+            else:
+                self.master.path_prompt(
+                    "Stream flows to: ",
+                    self.master.state.last_saveload,
+                    self.master.start_stream
+                )
         else:
             return urwid.ListBox.keypress(self, size, key)
-- 
cgit v1.2.3