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

Fixes #104
This commit is contained in:
Aldo Cortesi 2013-03-14 09:19:43 +13:00
parent c34d1e3de6
commit cfb5ba89ce
9 changed files with 48 additions and 18 deletions

View File

@ -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:
<pre class="terminal">
> mitmdump -s add_header.py

View File

@ -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

View File

@ -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('<body>', '<body><iframe src="%s" frameborder="0" height="0" width="0"></iframe>' % self._iframe_url)
if c > 0:
print 'Iframe injected!'
msg._ack()
msg.reply()
def main(argv):

View File

@ -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

View File

@ -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(

View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -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()