Introduce a filtered flow writer, and use it in dump.py

Fixes #104

3 files changed, 18 insertions, 5 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 4b5d1274..98904c84 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -426,7 +426,7 @@ class ConsoleMaster(flow.FlowMaster):
         path = os.path.expanduser(path)
         try:
             f = file(path, "wb")
-            flow.FlowMaster.start_stream(self, f)
+            flow.FlowMaster.start_stream(self, f, None)
         except IOError, v:
             return str(v)
         self.stream_path = path
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 3a315409..d716e433 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -93,7 +93,7 @@ class DumpMaster(flow.FlowMaster):
             path = os.path.expanduser(options.wfile)
             try:
                 f = file(path, "wb")
-                self.start_stream(f)
+                self.start_stream(f, self.filt)
             except IOError, v:
                 raise DumpError(v.strerror)
 
@@ -155,6 +155,7 @@ class DumpMaster(flow.FlowMaster):
         return "\n".join(" "*n + i for i in l)
 
     def _process_flow(self, f):
+        self.state.delete_flow(f)
         if self.filt and not f.match(self.filt):
             return
 
@@ -198,7 +199,6 @@ class DumpMaster(flow.FlowMaster):
             print >> self.outfile, "\n"
         if self.o.verbosity:
             self.outfile.flush()
-        self.state.delete_flow(f)
 
     def handle_log(self, l):
         self.add_event(l.msg)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 4fefee9f..13b32011 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -1588,8 +1588,8 @@ class FlowMaster(controller.Master):
                     self.stream.add(i)
             self.stop_stream()
 
-    def start_stream(self, fp):
-        self.stream = FlowWriter(fp)
+    def start_stream(self, fp, filt):
+        self.stream = FilteredFlowWriter(fp, filt)
 
     def stop_stream(self):
         self.stream.fo.close()
@@ -1635,3 +1635,16 @@ class FlowReader:
                 return
             raise FlowReadError("Invalid data format.")
 
+
+class FilteredFlowWriter:
+    def __init__(self, fo, filt):
+        self.fo = fo
+        self.filt = filt
+
+    def add(self, f):
+        if self.filt and not f.match(self.filt):
+            return
+        d = f._get_state()
+        tnetstring.dump(d, self.fo)
+
+