Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2025-02-03 16:55:33 +00:00
Code Issues Projects Releases Wiki Activity

Merge pull request #1706 from cortesi/taddons

Browse Source 
addons tests and bugfixes
This commit is contained in:
Aldo Cortesi 2016-11-03 16:33:25 +13:00 committed by GitHub
commit 50d393960c
13 changed files with 119 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mitmproxy/addons/__init__.py
View File

@ -1,7 +1,7 @@
from mitmproxy.addons import anticache
from mitmproxy.addons import anticomp
from mitmproxy.addons import clientplayback
from mitmproxy.addons import filestreamer
from mitmproxy.addons import streamfile
from mitmproxy.addons import onboarding
from mitmproxy.addons import replace
from mitmproxy.addons import script
@ -20,7 +20,7 @@ def default_addons():
stickyauth.StickyAuth(),
stickycookie.StickyCookie(),
script.ScriptLoader(),
filestreamer.FileStreamer(),
streamfile.StreamFile(),
streambodies.StreamBodies(),
replace.Replace(),
setheaders.SetHeaders(),

16
mitmproxy/addons/onboardingapp/app.py
View File

@ -6,7 +6,6 @@ import tornado.wsgi
from mitmproxy.utils import data
from mitmproxy.proxy import config
from mitmproxy.addons import wsgiapp
loader = tornado.template.Loader(data.pkg_data.path("addons/onboardingapp/templates"))
@ -92,18 +91,3 @@ application = tornado.web.Application(
],
# debug=True
)
class Onboarding(wsgiapp.WSGIApp):
def __init__(self):
super().__init__(Adapter(application), None, None)
self.enabled = False
def configure(self, options, updated):
self.host = options.app_host
self.port = options.app_port
self.enabled = options.app
def request(self, f):
if self.enabled:
super().request(f)

30
mitmproxy/addons/filestreamer.py → mitmproxy/addons/streamfile.py
View File

@ -5,9 +5,10 @@ from mitmproxy import flowfilter
from mitmproxy import io
class FileStreamer:
class StreamFile:
def __init__(self):
self.stream = None
self.filt = None
self.active_flows = set() # type: Set[flow.Flow]
def start_stream_to_path(self, path, mode, flt):
@ -15,29 +16,28 @@ class FileStreamer:
try:
f = open(path, mode)
except IOError as v:
return str(v)
raise exceptions.OptionsError(str(v))
self.stream = io.FilteredFlowWriter(f, flt)
self.active_flows = set()
def configure(self, options, updated):
# We're already streaming - stop the previous stream and restart
if self.stream:
self.done()
if options.outfile:
flt = None
if "filtstr" in updated:
if options.get("filtstr"):
flt = flowfilter.parse(options.filtstr)
if not flt:
self.filt = flowfilter.parse(options.filtstr)
if not self.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, flt)
if err:
raise exceptions.OptionsError(err)
if "streamfile" in updated:
if self.stream:
self.done()
if options.streamfile:
if options.streamfile_append:
mode = "ab"
else:
mode = "wb"
self.start_stream_to_path(options.streamfile, mode, self.filt)
def tcp_start(self, flow):
if self.stream:

6
mitmproxy/options.py
View File

@ -48,7 +48,8 @@ class Options(optmanager.OptManager):
stream_large_bodies: Optional[int] = None,
verbosity: int = 2,
default_contentview: str = "auto",
outfile: Optional[Tuple[str, str]] = None,
streamfile: Optional[str] = None,
streamfile_append: bool = False,
server_replay_ignore_content: bool = False,
server_replay_ignore_params: Sequence[str] = (),
server_replay_ignore_payload_params: Sequence[str] = (),
@ -108,7 +109,8 @@ class Options(optmanager.OptManager):
self.stream_large_bodies = stream_large_bodies
self.verbosity = verbosity
self.default_contentview = default_contentview
self.outfile = outfile
self.streamfile = streamfile
self.streamfile_append = streamfile_append
self.server_replay_ignore_content = server_replay_ignore_content
self.server_replay_ignore_params = server_replay_ignore_params
self.server_replay_ignore_payload_params = server_replay_ignore_payload_params

18
mitmproxy/test/taddons.py
View File

@ -1,7 +1,10 @@
import contextlib
import mitmproxy.master
import mitmproxy.options
from mitmproxy import proxy
from mitmproxy import events
from mitmproxy import exceptions
class RecordingMaster(mitmproxy.master.Master):
@ -36,6 +39,15 @@ class context:
self.wrapped = None
return False
@contextlib.contextmanager
def _rollback(self, opts, updates):
old = opts._opts.copy()
try:
yield
except exceptions.OptionsError as e:
opts.__dict__["_opts"] = old
raise
def cycle(self, addon, f):
"""
Cycles the flow through the events for the flow. Stops if a reply
@ -55,6 +67,6 @@ class context:
Options object with the given keyword arguments, then calls the
configure method on the addon with the updated value.
"""
for k, v in kwargs.items():
setattr(self.options, k, v)
addon.configure(self.options, kwargs.keys())
with self._rollback(self.options, kwargs):
self.options.update(**kwargs)
addon.configure(self.options, kwargs.keys())

17
mitmproxy/tools/cmdline.py
View File

@ -140,8 +140,8 @@ def get_common_options(args):
raise exceptions.OptionsError(e)
setheaders.append(p)
if args.outfile and args.outfile[0] == args.rfile:
if args.outfile[1] == "wb":
if args.streamfile and args.streamfile[0] == args.rfile:
if args.streamfile[1] == "wb":
raise exceptions.OptionsError(
"Cannot use '{}' for both reading and writing flows. "
"Are you looking for --afile?".format(args.rfile)
@ -228,7 +228,8 @@ def get_common_options(args):
stickyauth=stickyauth,
stream_large_bodies=stream_large_bodies,
showhost=args.showhost,
outfile=args.outfile,
streamfile=args.streamfile[0] if args.streamfile else None,
streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False,
verbosity=args.verbose,
server_replay_nopop=args.server_replay_nopop,
server_replay_ignore_content=args.server_replay_ignore_content,
@ -339,15 +340,15 @@ def basic_options(parser):
action="store_const", dest="verbose", default=2, const=3,
help="Increase log verbosity."
)
outfile = parser.add_mutually_exclusive_group()
outfile.add_argument(
streamfile = parser.add_mutually_exclusive_group()
streamfile.add_argument(
"-w", "--wfile",
action="store", dest="outfile", type=lambda f: (f, "wb"),
action="store", dest="streamfile", type=lambda f: (f, "w"),
help="Write flows to file."
)
outfile.add_argument(
streamfile.add_argument(
"-a", "--afile",
action="store", dest="outfile", type=lambda f: (f, "ab"),
action="store", dest="streamfile", type=lambda f: (f, "a"),
help="Append flows to file."
)
parser.add_argument(

6
mitmproxy/tools/console/flowlist.py
View File

@ -393,13 +393,13 @@ class FlowListBox(urwid.ListBox):
val = not self.master.options.order_reversed
self.master.options.order_reversed = val
elif key == "W":
if self.master.options.outfile:
self.master.options.outfile = None
if self.master.options.streamfile:
self.master.options.streamfile = None
else:
signals.status_prompt_path.send(
self,
prompt="Stream flows to",
callback= lambda path: self.master.options.update(outfile=(path, "ab"))
callback= lambda path: self.master.options.update(streamfile=path)
)
else:
return urwid.ListBox.keypress(self, size, key)

4
mitmproxy/tools/console/statusbar.py
View File

@ -223,8 +223,8 @@ class StatusBar(urwid.WidgetWrap):
r.append(("heading_key", "s"))
r.append("cripts:%s]" % len(self.master.options.scripts))
if self.master.options.outfile:
r.append("[W:%s]" % self.master.options.outfile[0])
if self.master.options.streamfile:
r.append("[W:%s]" % self.master.options.streamfile)
return r

9
mitmproxy/tools/web/master.py
View File

@ -111,15 +111,6 @@ class WebMaster(master.Master):
"error"
)
if options.outfile:
err = self.start_stream_to_path(
options.outfile[0],
options.outfile[1]
)
if err:
print("Stream file error: {}".format(err), file=sys.stderr)
sys.exit(1)
def _sig_add(self, view, flow):
app.ClientConnection.broadcast(
type="UPDATE_FLOWS",

44
test/mitmproxy/addons/test_filestreamer.py
View File

@ -1,44 +0,0 @@
from mitmproxy.test import tflow
from mitmproxy.test import tutils
from .. import mastertest
import os.path
from mitmproxy.addons import filestreamer
from mitmproxy import master
from mitmproxy import io
from mitmproxy import options
from mitmproxy import proxy
class TestStream(mastertest.MasterTest):
def test_stream(self):
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
def r():
r = io.FlowReader(open(p, "rb"))
return list(r.stream())
o = options.Options(
outfile = (p, "wb")
)
m = master.Master(o, proxy.DummyServer())
sa = filestreamer.FileStreamer()
m.addons.add(sa)
f = tflow.tflow(resp=True)
m.request(f)
m.response(f)
m.addons.remove(sa)
assert r()[0].response
m.options.outfile = (p, "ab")
m.addons.add(sa)
f = tflow.tflow()
m.request(f)
m.addons.remove(sa)
assert not r()[1].response

5
test/mitmproxy/addons/test_intercept.py
View File

@ -36,3 +36,8 @@ def test_simple():
f = tflow.tflow(resp=False)
tctx.cycle(r, f)
assert not f.intercepted
f = tflow.tflow(resp=True)
f.reply._state = "handled"
r.response(f)
assert f.intercepted

64
test/mitmproxy/addons/test_streamfile.py Normal file
View File

@ -0,0 +1,64 @@
from mitmproxy.test import tflow
from mitmproxy.test import tutils
from mitmproxy.test import taddons
import os.path
from mitmproxy import io
from mitmproxy import exceptions
from mitmproxy.tools import dump
from mitmproxy.addons import streamfile
def test_configure():
sa = streamfile.StreamFile()
with taddons.context(options=dump.Options()) as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tutils.raises(
exceptions.OptionsError,
tctx.configure, sa, streamfile=tdir
)
tutils.raises(
"invalid filter",
tctx.configure, sa, streamfile=p, filtstr="~~"
)
def rd(p):
x = io.FlowReader(open(p, "rb"))
return list(x.stream())
def test_tcp():
sa = streamfile.StreamFile()
with taddons.context() as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tctx.configure(sa, streamfile=p)
tt = tflow.ttcpflow()
sa.tcp_start(tt)
sa.tcp_end(tt)
tctx.configure(sa, streamfile=None)
assert rd(p)
def test_simple():
sa = streamfile.StreamFile()
with taddons.context() as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tctx.configure(sa, streamfile=p)
f = tflow.tflow(resp=True)
sa.request(f)
sa.response(f)
tctx.configure(sa, streamfile=None)
assert rd(p)[0].response
tctx.configure(sa, streamfile=p, streamfile_append=True)
f = tflow.tflow()
sa.request(f)
tctx.configure(sa, streamfile=None)
assert not rd(p)[1].response

28
test/mitmproxy/test_dump.py
View File

@ -2,7 +2,6 @@ from mitmproxy.test import tflow
import os
import io
import mitmproxy.io
from mitmproxy.tools import dump
from mitmproxy import exceptions
from mitmproxy import proxy
@ -126,33 +125,6 @@ class TestDumpMaster(mastertest.MasterTest):
f = self.cycle(m, b"content")
assert f.request.headers["one"] == "two"
def test_write(self):
with tutils.tmpdir() as d:
p = os.path.join(d, "a")
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "wb"), verbosity=0), 1, b""
)
assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 1
def test_write_append(self):
with tutils.tmpdir() as d:
p = os.path.join(d, "a.append")
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "wb"), verbosity=0),
1, b""
)
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "ab"), verbosity=0),
1, b""
)
assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 2
def test_write_err(self):
tutils.raises(
exceptions.OptionsError,
self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb")
)
def test_script(self):
ret = self.dummy_cycle(
self.mkmaster(