From cfb5ba89ce96594b2f8d51f27c9b2ee41ecf18e5 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Mar 2013 09:19:43 +1300 Subject: [PATCH] Introduce a filtered flow writer, and use it in dump.py Fixes #104 --- doc-src/scripting/inlinescripts.html | 2 +- examples/flowbasic | 8 ++++---- examples/iframe_injector | 4 ++-- examples/proxapp | 4 ++-- examples/stickycookies | 4 ++-- libmproxy/console/__init__.py | 2 +- libmproxy/dump.py | 4 ++-- libmproxy/flow.py | 17 +++++++++++++++-- test/test_flow.py | 21 +++++++++++++++++++-- 9 files changed, 48 insertions(+), 18 deletions(-) diff --git a/doc-src/scripting/inlinescripts.html b/doc-src/scripting/inlinescripts.html index 39752b7f0..2d53df80c 100644 --- a/doc-src/scripting/inlinescripts.html +++ b/doc-src/scripting/inlinescripts.html @@ -13,7 +13,7 @@ lets the script interact with the global mitmproxy state. The __response__ event also gets an instance of Flow, which we can use to manipulate the response itself. -We can now run this script as follows: +We can now run this script using mitmdump or mitmproxy as follows: 
 > mitmdump -s add_header.py
diff --git a/examples/flowbasic b/examples/flowbasic
index b419f9e75..b8184262f 100755
--- a/examples/flowbasic
+++ b/examples/flowbasic
@@ -3,8 +3,8 @@
     This example shows how to build a proxy based on mitmproxy's Flow
     primitives.
 
-    Note that request and response messages are not automatically acked, so we
-    need to implement handlers to do this.
+    Note that request and response messages are not automatically replied to,
+    so we need to implement handlers to do this.
 """
 import os
 from libmproxy import proxy, flow
@@ -19,13 +19,13 @@ class MyMaster(flow.FlowMaster):
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
         if f:
-            r._ack()
+            r.reply()
         return f
 
     def handle_response(self, r):
         f = flow.FlowMaster.handle_response(self, r)
         if f:
-            r._ack()
+            r.reply()
         print f
         return f
 
diff --git a/examples/iframe_injector b/examples/iframe_injector
index 6dd28674c..8b1e02f15 100755
--- a/examples/iframe_injector
+++ b/examples/iframe_injector
@@ -23,14 +23,14 @@ class InjectingMaster(controller.Master):
     def handle_request(self, msg):
         if 'Accept-Encoding' in msg.headers:
             msg.headers["Accept-Encoding"] = 'none'
-        msg._ack()
+        msg.reply()
 
     def handle_response(self, msg):
         if msg.content:
             c = msg.replace('', '' % self._iframe_url)
             if c > 0:
                 print 'Iframe injected!'
-        msg._ack()
+        msg.reply()
 
 
 def main(argv):
diff --git a/examples/proxapp b/examples/proxapp
index e8179528a..eb5bdbb70 100755
--- a/examples/proxapp
+++ b/examples/proxapp
@@ -23,13 +23,13 @@ class MyMaster(flow.FlowMaster):
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
         if f:
-            r._ack()
+            r.reply()
         return f
 
     def handle_response(self, r):
         f = flow.FlowMaster.handle_response(self, r)
         if f:
-            r._ack()
+            r.reply()
         print f
         return f
 
diff --git a/examples/stickycookies b/examples/stickycookies
index 88bf00638..b07820fca 100755
--- a/examples/stickycookies
+++ b/examples/stickycookies
@@ -25,13 +25,13 @@ class StickyMaster(controller.Master):
             self.stickyhosts[hid] = msg.headers["cookie"]
         elif hid in self.stickyhosts:
             msg.headers["cookie"] = self.stickyhosts[hid]
-        msg._ack()
+        msg.reply()
 
     def handle_response(self, msg):
         hid = (msg.request.host, msg.request.port)
         if msg.headers["set-cookie"]:
             self.stickyhosts[hid] = msg.headers["set-cookie"]
-        msg._ack()
+        msg.reply()
 
 
 config = proxy.ProxyConfig(
diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 4b5d12743..98904c84b 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 3a315409b..d716e4338 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 4fefee9f1..13b320111 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)
+
+
diff --git a/test/test_flow.py b/test/test_flow.py
index 6aa898adc..c1ae1a9f8 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -497,6 +497,23 @@ class TestSerialize:
         fm = flow.FlowMaster(None, s)
         fm.load_flows(r)
         assert len(s._flow_list) == 6
+        
+    def test_filter(self):
+        sio = StringIO()
+        fl = filt.parse("~c 200")
+        w = flow.FilteredFlowWriter(sio, fl)
+
+        f = tutils.tflow_full()
+        f.response.code = 200
+        w.add(f)
+
+        f = tutils.tflow_full()
+        f.response.code = 201
+        w.add(f)
+
+        sio.seek(0)
+        r = flow.FlowReader(sio)
+        assert len(list(r.stream()))
 
 
     def test_error(self):
@@ -723,7 +740,7 @@ class TestFlowMaster:
             fm = flow.FlowMaster(None, s)
             tf = tutils.tflow_full()
 
-            fm.start_stream(file(p, "ab"))
+            fm.start_stream(file(p, "ab"), None)
             fm.handle_request(tf.request)
             fm.handle_response(tf.response)
             fm.stop_stream()
@@ -731,7 +748,7 @@ class TestFlowMaster:
             assert r()[0].response
 
             tf = tutils.tflow_full()
-            fm.start_stream(file(p, "ab"))
+            fm.start_stream(file(p, "ab"), None)
             fm.handle_request(tf.request)
             fm.shutdown()