Add streaming to FlowMaster

2 files changed, 24 insertions, 8 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 72776542..4034c188 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -928,7 +928,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.state.killall(self)
         if self.stream:
             self.stream.fo.close()
-        controller.Master.shutdown(self)
+        flow.FlowMaster.shutdown(self)
 
     def sync_list_view(self):
         self.flow_list_walker._modified()
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 5f5cad4c..17e88bc1 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -17,7 +17,7 @@
     This module provides more sophisticated flow tracking. These match requests
     with their responses, and provide filtering and interception facilities.
 """
-import hashlib, Cookie, cookielib, copy, re, urlparse
+import hashlib, Cookie, cookielib, copy, re, urlparse, os
 import time
 import tnetstring, filt, script, utils, encoding, proxy
 from email.utils import parsedate_tz, formatdate, mktime_tz
@@ -1200,6 +1200,8 @@ class FlowMaster(controller.Master):
         self.refresh_server_playback = False
         self.replacehooks = ReplaceHooks()
 
+        self.stream = None
+
     def add_event(self, e, level="info"):
         """
             level: info, error
@@ -1415,21 +1417,35 @@ class FlowMaster(controller.Master):
 
     def handle_response(self, r):
         f = self.state.add_response(r)
-        self.replacehooks.run(f)
         if f:
+            self.replacehooks.run(f)
             self.run_script_hook("response", f)
-        if self.client_playback:
-            self.client_playback.clear(f)
-        if not f:
-            r._ack()
-        if f:
+            if self.client_playback:
+                self.client_playback.clear(f)
             self.process_new_response(f)
+            if self.stream:
+                self.stream.add(f)
+        else:
+            r._ack()
         return f
 
     def shutdown(self):
         if self.script:
             self.load_script(None)
         controller.Master.shutdown(self)
+        if self.stream:
+            for i in self.state._flow_list:
+                if not i.response:
+                    self.stream.add(i)
+            self.stream.fo.close()
+            self.stop_stream()
+
+    def start_stream(self, fp):
+        self.stream = FlowWriter(fp)
+
+    def stop_stream(self):
+        self.stream = None
+
 
 
 class FlowWriter: