Some work on stream-to-file

- stream -> filestreamer throughout
- move active_flows set into FileStreamer addon

mitmproxy/builtins/__init__.py

@@ -2,10 +2,10 @@ from __future__ import absolute_import, print_function, division
 
 from mitmproxy.builtins import anticache
 from mitmproxy.builtins import anticomp
+from mitmproxy.builtins import filestreamer
 from mitmproxy.builtins import stickyauth
 from mitmproxy.builtins import stickycookie
 from mitmproxy.builtins import script
-from mitmproxy.builtins import stream
 from mitmproxy.builtins import replace
 
 
@@ -16,6 +16,6 @@ def default_addons():
         stickyauth.StickyAuth(),
         stickycookie.StickyCookie(),
         script.ScriptLoader(),
-        stream.Stream(),
+        filestreamer.FileStreamer(),
         replace.Replace(),
     ]

mitmproxy/builtins/filestreamer.py

@@ -1,14 +1,14 @@
 from __future__ import absolute_import, print_function, division
 import os.path
 
-from mitmproxy import ctx
 from mitmproxy import exceptions
 from mitmproxy.flow import io
 
 
-class Stream:
+class FileStreamer:
     def __init__(self):
         self.stream = None
+        self.active_flows = set()  # type: Set[models.Flow]
 
     def start_stream_to_path(self, path, mode, filt):
         path = os.path.expanduser(path)
@@ -17,6 +17,7 @@ class Stream:
         except IOError as v:
             return str(v)
         self.stream = io.FilteredFlowWriter(f, filt)
+        self.active_flows = set()
 
     def configure(self, options):
         # We're already streaming - stop the previous stream and restart
@@ -38,17 +39,28 @@ class Stream:
             if err:
                 raise exceptions.OptionsError(err)
 
-    def done(self):
+    def tcp_open(self, flow):
         if self.stream:
-            for flow in ctx.master.active_flows:
-                self.stream.add(flow)
-            self.stream.fo.close()
-            self.stream = None
+            self.active_flows.add(flow)
 
     def tcp_close(self, flow):
         if self.stream:
             self.stream.add(flow)
+            self.active_flows.discard(flow)
 
     def response(self, flow):
         if self.stream:
             self.stream.add(flow)
+            self.active_flows.discard(flow)
+
+    def request(self, flow):
+        if self.stream:
+            self.active_flows.add(flow)
+
+    def done(self):
+        if self.stream:
+            for flow in self.active_flows:
+                self.stream.add(flow)
+            self.active_flows = set([])
+            self.stream.fo.close()
+            self.stream = None

mitmproxy/flow/master.py

@@ -30,7 +30,6 @@ class FlowMaster(controller.Master):
         if server:
             self.add_server(server)
         self.state = state
-        self.active_flows = set()  # type: Set[models.Flow]
         self.server_playback = None  # type: Optional[modules.ServerPlaybackState]
         self.client_playback = None  # type: Optional[modules.ClientPlaybackState]
         self.kill_nonreplay = False
@@ -329,7 +328,6 @@ class FlowMaster(controller.Master):
                 return
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
-        self.active_flows.add(f)
         if not f.reply.acked:
             self.setheaders.run(f)
         if not f.reply.acked:
@@ -348,7 +346,6 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def response(self, f):
-        self.active_flows.discard(f)
         self.state.update_flow(f)
         if not f.reply.acked:
             self.setheaders.run(f)
@@ -367,7 +364,7 @@ class FlowMaster(controller.Master):
     def tcp_open(self, flow):
         # TODO: This would break mitmproxy currently.
         # self.state.add_flow(flow)
-        self.active_flows.add(flow)
+        pass
 
     @controller.handler
     def tcp_message(self, flow):
@@ -382,4 +379,4 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def tcp_close(self, flow):
-        self.active_flows.discard(flow)
+        pass

test/mitmproxy/builtins/test_filestreamer.py

@@ -4,7 +4,7 @@ from .. import tutils, mastertest
 
 import os.path
 
-from mitmproxy.builtins import stream
+from mitmproxy.builtins import filestreamer
 from mitmproxy.flow import master, FlowReader
 from mitmproxy.flow import state
 from mitmproxy.flow import options
@@ -27,7 +27,7 @@ class TestStream(mastertest.MasterTest):
                 None,
                 s
             )
-            sa = stream.Stream()
+            sa = filestreamer.FileStreamer()
 
             m.addons.add(sa)
             f = tutils.tflow(resp=True)