Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 08:11:00 +00:00
Code Issues Projects Releases Wiki Activity

Merge pull request #1353 from cortesi/stream

Browse Source 
Streaming to file -> addon
This commit is contained in:
Aldo Cortesi 2016-07-14 14:09:59 +12:00 committed by GitHub
commit b8a23eeaa3
8 changed files with 167 additions and 118 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/builtins/__init__.py
View File

@ -4,6 +4,7 @@ from mitmproxy.builtins import anticache
from mitmproxy.builtins import anticomp
from mitmproxy.builtins import stickyauth
from mitmproxy.builtins import stickycookie
from mitmproxy.builtins import stream
def default_addons():
@ -12,4 +13,5 @@ def default_addons():
anticomp.AntiComp(),
stickyauth.StickyAuth(),
stickycookie.StickyCookie(),
stream.Stream(),
]

54
mitmproxy/builtins/stream.py Normal file
View File

@ -0,0 +1,54 @@
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:
def __init__(self):
self.stream = None
def start_stream_to_path(self, path, mode, filt):
path = os.path.expanduser(path)
try:
f = open(path, mode)
except IOError as v:
return str(v)
self.stream = io.FilteredFlowWriter(f, filt)
def configure(self, options):
# We're already streaming - stop the previous stream and restart
if self.stream:
self.done()
if options.outfile:
filt = None
if options.get("filtstr"):
filt = filt.parse(options.filtstr)
if not filt:
raise exceptions.OptionsError(
"Invalid filter specification: %s" % options.filtstr
)
path, mode = options.outfile
if mode not in ("wb", "ab"):
raise exceptions.OptionsError("Invalid mode.")
err = self.start_stream_to_path(path, mode, filt)
if err:
raise exceptions.OptionsError(err)
def done(self):
if self.stream:
for flow in ctx.master.active_flows:
self.stream.add(flow)
self.stream.fo.close()
self.stream = None
def tcp_close(self, flow):
if self.stream:
self.stream.add(flow)
def response(self, flow):
if self.stream:
self.stream.add(flow)

22
mitmproxy/dump.py
View File

@ -52,16 +52,17 @@ class Options(options.Options):
"replay_ignore_content",
"replay_ignore_params",
"replay_ignore_payload_params",
"replay_ignore_host"
"replay_ignore_host",
"tfile"
]
class DumpMaster(flow.FlowMaster):
def __init__(self, server, options, outfile=None):
def __init__(self, server, options):
flow.FlowMaster.__init__(self, options, server, flow.State())
self.addons.add(*builtins.default_addons())
self.outfile = outfile
self.o = options
self.showhost = options.showhost
self.replay_ignore_params = options.replay_ignore_params
@ -82,15 +83,6 @@ class DumpMaster(flow.FlowMaster):
else:
self.filt = None
if options.outfile:
err = self.start_stream_to_path(
options.outfile[0],
options.outfile[1],
self.filt
)
if err:
raise DumpError(err)
if options.replacements:
for i in options.replacements:
self.replacehooks.add(*i)
@ -163,7 +155,7 @@ class DumpMaster(flow.FlowMaster):
def echo(self, text, indent=None, **style):
if indent:
text = self.indent(indent, text)
click.secho(text, file=self.outfile, **style)
click.secho(text, file=self.options.tfile, **style)
def _echo_message(self, message):
if self.options.flow_detail >= 2 and hasattr(message, "headers"):
@ -312,8 +304,8 @@ class DumpMaster(flow.FlowMaster):
if f.error:
self.echo(" << {}".format(f.error.msg), bold=True, fg="red")
if self.outfile:
self.outfile.flush()
if self.options.tfile:
self.options.tfile.flush()
def _process_flow(self, f):
if self.filt and not f.match(self.filt):

28
mitmproxy/flow/master.py
View File

@ -46,7 +46,6 @@ class FlowMaster(controller.Master):
self.replay_ignore_content = None
self.replay_ignore_host = False
self.stream = None
self.apps = modules.AppRegistry()
def start_app(self, host, port):
@ -409,8 +408,6 @@ class FlowMaster(controller.Master):
if not f.reply.acked:
if self.client_playback:
self.client_playback.clear(f)
if self.stream:
self.stream.add(f)
return f
def handle_intercept(self, f):
@ -471,33 +468,8 @@ class FlowMaster(controller.Master):
@controller.handler
def tcp_close(self, flow):
self.active_flows.discard(flow)
if self.stream:
self.stream.add(flow)
self.run_scripts("tcp_close", flow)
def shutdown(self):
super(FlowMaster, self).shutdown()
# Add all flows that are still active
if self.stream:
for flow in self.active_flows:
self.stream.add(flow)
self.stop_stream()
self.unload_scripts()
def start_stream(self, fp, filt):
self.stream = io.FilteredFlowWriter(fp, filt)
def stop_stream(self):
self.stream.fo.close()
self.stream = None
def start_stream_to_path(self, path, mode="wb", filt=None):
path = os.path.expanduser(path)
try:
f = open(path, mode)
self.start_stream(f, filt)
except IOError as v:
return str(v)
self.stream_path = path

46
test/mitmproxy/builtins/test_stream.py Normal file
View File

@ -0,0 +1,46 @@
from __future__ import absolute_import, print_function, division
from .. import tutils, mastertest
import os.path
from mitmproxy.builtins import stream
from mitmproxy.flow import master, FlowReader
from mitmproxy.flow import state
from mitmproxy import options
class TestStream(mastertest.MasterTest):
def test_stream(self):
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
def r():
r = FlowReader(open(p, "rb"))
return list(r.stream())
s = state.State()
m = master.FlowMaster(
options.Options(
outfile = (p, "wb")
),
None,
s
)
sa = stream.Stream()
m.addons.add(sa)
f = tutils.tflow(resp=True)
self.invoke(m, "request", f)
self.invoke(m, "response", f)
m.addons.remove(sa)
assert r()[0].response
m.options.outfile = (p, "ab")
m.addons.add(sa)
f = tutils.tflow()
self.invoke(m, "request", f)
m.addons.remove(sa)
assert not r()[1].response

13
test/mitmproxy/mastertest.py
View File

@ -18,15 +18,14 @@ class MasterTest:
l = proxy.Log("connect")
l.reply = mock.MagicMock()
master.log(l)
master.clientconnect(f.client_conn)
master.serverconnect(f.server_conn)
master.request(f)
self.invoke(master, "clientconnect", f.client_conn)
self.invoke(master, "clientconnect", f.client_conn)
self.invoke(master, "serverconnect", f.server_conn)
self.invoke(master, "request", f)
if not f.error:
f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content))
f.reply.acked = False
f = master.response(f)
f.client_conn.reply.acked = False
master.clientdisconnect(f.client_conn)
self.invoke(master, "response", f)
self.invoke(master, "clientdisconnect", f)
return f
def dummy_cycle(self, master, n, content):

89
test/mitmproxy/test_dump.py
View File

@ -4,31 +4,33 @@ from mitmproxy.exceptions import ContentViewException
import netlib.tutils
from mitmproxy import dump, flow, models
from mitmproxy import dump, flow, models, exceptions
from . import tutils, mastertest
import mock
def test_strfuncs():
o = dump.Options()
o = dump.Options(
tfile = StringIO(),
flow_detail = 0,
)
m = dump.DumpMaster(None, o)
m.outfile = StringIO()
m.o.flow_detail = 0
m.echo_flow(tutils.tflow())
assert not m.outfile.getvalue()
assert not o.tfile.getvalue()
m.o.flow_detail = 4
m.echo_flow(tutils.tflow())
assert m.outfile.getvalue()
assert o.tfile.getvalue()
m.outfile = StringIO()
o.tfile = StringIO()
m.echo_flow(tutils.tflow(resp=True))
assert "<<" in m.outfile.getvalue()
assert "<<" in o.tfile.getvalue()
m.outfile = StringIO()
o.tfile = StringIO()
m.echo_flow(tutils.tflow(err=True))
assert "<<" in m.outfile.getvalue()
assert "<<" in o.tfile.getvalue()
flow = tutils.tflow()
flow.request = netlib.tutils.treq()
@ -50,25 +52,32 @@ def test_strfuncs():
def test_contentview(get_content_view):
get_content_view.side_effect = ContentViewException(""), ("x", iter([]))
o = dump.Options(flow_detail=4, verbosity=3)
m = dump.DumpMaster(None, o, StringIO())
o = dump.Options(
flow_detail=4,
verbosity=3,
tfile=StringIO(),
)
m = dump.DumpMaster(None, o)
m.echo_flow(tutils.tflow())
assert "Content viewer failed" in m.outfile.getvalue()
assert "Content viewer failed" in m.options.tfile.getvalue()
class TestDumpMaster(mastertest.MasterTest):
def dummy_cycle(self, master, n, content):
mastertest.MasterTest.dummy_cycle(self, master, n, content)
return master.outfile.getvalue()
return master.options.tfile.getvalue()
def mkmaster(self, filt, **options):
cs = StringIO()
if "verbosity" not in options:
options["verbosity"] = 0
if "flow_detail" not in options:
options["flow_detail"] = 0
o = dump.Options(filtstr=filt, **options)
return dump.DumpMaster(None, o, outfile=cs)
o = dump.Options(
filtstr=filt,
tfile=StringIO(),
**options
)
return dump.DumpMaster(None, o)
def test_basic(self):
for i in (1, 2, 3):
@ -89,31 +98,33 @@ class TestDumpMaster(mastertest.MasterTest):
)
def test_error(self):
cs = StringIO()
o = dump.Options(flow_detail=1)
m = dump.DumpMaster(None, o, outfile=cs)
o = dump.Options(
tfile=StringIO(),
flow_detail=1
)
m = dump.DumpMaster(None, o)
f = tutils.tflow(err=True)
m.request(f)
assert m.error(f)
assert "error" in cs.getvalue()
assert "error" in o.tfile.getvalue()
def test_missing_content(self):
cs = StringIO()
o = dump.Options(flow_detail=3)
m = dump.DumpMaster(None, o, outfile=cs)
o = dump.Options(
flow_detail=3,
tfile=StringIO(),
)
m = dump.DumpMaster(None, o)
f = tutils.tflow()
f.request.content = None
m.request(f)
f.response = models.HTTPResponse.wrap(netlib.tutils.tresp())
f.response.content = None
m.response(f)
assert "content missing" in cs.getvalue()
assert "content missing" in o.tfile.getvalue()
def test_replay(self):
cs = StringIO()
o = dump.Options(server_replay=["nonexistent"], kill=True)
tutils.raises(dump.DumpError, dump.DumpMaster, None, o, outfile=cs)
tutils.raises(dump.DumpError, dump.DumpMaster, None, o)
with tutils.tmpdir() as t:
p = os.path.join(t, "rep")
@ -122,7 +133,7 @@ class TestDumpMaster(mastertest.MasterTest):
o = dump.Options(server_replay=[p], kill=True)
o.verbosity = 0
o.flow_detail = 0
m = dump.DumpMaster(None, o, outfile=cs)
m = dump.DumpMaster(None, o)
self.cycle(m, b"content")
self.cycle(m, b"content")
@ -130,13 +141,13 @@ class TestDumpMaster(mastertest.MasterTest):
o = dump.Options(server_replay=[p], kill=False)
o.verbosity = 0
o.flow_detail = 0
m = dump.DumpMaster(None, o, outfile=cs)
m = dump.DumpMaster(None, o)
self.cycle(m, b"nonexistent")
o = dump.Options(client_replay=[p], kill=False)
o.verbosity = 0
o.flow_detail = 0
m = dump.DumpMaster(None, o, outfile=cs)
m = dump.DumpMaster(None, o)
def test_read(self):
with tutils.tmpdir() as t:
@ -172,20 +183,24 @@ class TestDumpMaster(mastertest.MasterTest):
assert len(m.apps.apps) == 1
def test_replacements(self):
cs = StringIO()
o = dump.Options(replacements=[(".*", "content", "foo")])
o = dump.Options(
replacements=[(".*", "content", "foo")],
tfile = StringIO(),
)
o.verbosity = 0
o.flow_detail = 0
m = dump.DumpMaster(None, o, outfile=cs)
m = dump.DumpMaster(None, o)
f = self.cycle(m, b"content")
assert f.request.content == b"foo"
def test_setheader(self):
cs = StringIO()
o = dump.Options(setheaders=[(".*", "one", "two")])
o = dump.Options(
setheaders=[(".*", "one", "two")],
tfile=StringIO()
)
o.verbosity = 0
o.flow_detail = 0
m = dump.DumpMaster(None, o, outfile=cs)
m = dump.DumpMaster(None, o)
f = self.cycle(m, b"content")
assert f.request.headers["one"] == "two"
@ -212,7 +227,7 @@ class TestDumpMaster(mastertest.MasterTest):
def test_write_err(self):
tutils.raises(
dump.DumpError,
exceptions.OptionsError,
self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb")
)

31
test/mitmproxy/test_flow.py
View File

@ -1,5 +1,3 @@
import os.path
import mock
import io
@ -887,35 +885,6 @@ class TestFlowMaster:
fm.process_new_request(f)
assert "killed" in f.error.msg
def test_stream(self):
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
def read():
with open(p, "rb") as f:
r = flow.FlowReader(f)
return list(r.stream())
s = flow.State()
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow(resp=True)
with open(p, "ab") as tmpfile:
fm.start_stream(tmpfile, None)
fm.request(f)
fm.response(f)
fm.stop_stream()
assert read()[0].response
with open(p, "ab") as tmpfile:
f = tutils.tflow()
fm.start_stream(tmpfile, None)
fm.request(f)
fm.shutdown()
assert not read()[1].response
class TestRequest: