From cfb5ba89ce96594b2f8d51f27c9b2ee41ecf18e5 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Thu, 14 Mar 2013 09:19:43 +1300
Subject: Introduce a filtered flow writer, and use it in dump.py

Fixes #104
---
 libmproxy/console/__init__.py |  2 +-
 libmproxy/dump.py             |  4 ++--
 libmproxy/flow.py             | 17 +++++++++++++++--
 3 files changed, 18 insertions(+), 5 deletions(-)

(limited to 'libmproxy')

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)
+
+
-- 
cgit v1.2.3