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 #1348 from cortesi/addons

Browse Source 
Addons
This commit is contained in:
Aldo Cortesi 2016-07-14 00:42:12 +12:00 committed by GitHub
commit 07c3f90c5b
18 changed files with 194 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
mitmproxy/addons.py Normal file
View File

@ -0,0 +1,63 @@
from __future__ import absolute_import, print_function, division
from mitmproxy import exceptions
import pprint
def _get_name(itm):
return getattr(itm, "name", itm.__class__.__name__)
class Addons(object):
def __init__(self, master):
self.chain = []
self.master = master
master.options.changed.connect(self.options_update)
def options_update(self, options):
for i in self.chain:
with self.master.handlecontext():
i.configure(options)
def add(self, *addons):
self.chain.extend(addons)
for i in addons:
self.invoke_with_context(i, "configure", self.master.options)
def remove(self, addon):
self.chain = [i for i in self.chain if i is not addon]
self.invoke_with_context(addon, "done")
def done(self):
for i in self.chain:
self.invoke_with_context(i, "done")
def has_addon(self, name):
"""
Is an addon with this name registered?
"""
for i in self.chain:
if _get_name(i) == name:
return True
def __len__(self):
return len(self.chain)
def __str__(self):
return pprint.pformat([str(i) for i in self.chain])
def invoke_with_context(self, addon, name, *args, **kwargs):
with self.master.handlecontext():
self.invoke(addon, name, *args, **kwargs)
def invoke(self, addon, name, *args, **kwargs):
func = getattr(addon, name, None)
if func:
if not callable(func):
raise exceptions.AddonError(
"Addon handler %s not callable" % name
)
func(*args, **kwargs)
def __call__(self, name, *args, **kwargs):
for i in self.chain:
self.invoke(i, name, *args, **kwargs)

9
mitmproxy/builtins/__init__.py Normal file
View File

@ -0,0 +1,9 @@
from __future__ import absolute_import, print_function, division
from mitmproxy.builtins import anticomp
def default_addons():
return [
anticomp.AntiComp(),
]

13
mitmproxy/builtins/anticomp.py Normal file
View File

@ -0,0 +1,13 @@
from __future__ import absolute_import, print_function, division
class AntiComp:
def __init__(self):
self.enabled = False
def configure(self, options):
self.enabled = options.anticomp
def request(self, flow):
if self.enabled:
flow.request.anticomp()

6
mitmproxy/console/master.py
View File

@ -15,6 +15,7 @@ import weakref
import urwid
from mitmproxy import builtins
from mitmproxy import contentviews
from mitmproxy import controller
from mitmproxy import exceptions
@ -216,10 +217,10 @@ class ConsoleMaster(flow.FlowMaster):
palette = []
def __init__(self, server, options):
flow.FlowMaster.__init__(self, server, ConsoleState())
flow.FlowMaster.__init__(self, options, server, ConsoleState())
self.addons.add(*builtins.default_addons())
self.stream_path = None
self.options = options
self.options.errored.connect(self.options_error)
if options.replacements:
@ -252,7 +253,6 @@ class ConsoleMaster(flow.FlowMaster):
self.refresh_server_playback = options.refresh_server_playback
self.anticache = options.anticache
self.anticomp = options.anticomp
self.killextra = options.kill
self.rheaders = options.rheaders
self.nopop = options.nopop

4
mitmproxy/console/options.py
View File

@ -102,7 +102,7 @@ class Options(urwid.WidgetWrap):
select.Option(
"Anti-Compression",
"o",
lambda: master.anticomp,
lambda: master.options.anticomp,
self.toggle_anticomp
),
select.Option(
@ -177,7 +177,7 @@ class Options(urwid.WidgetWrap):
self.master.anticache = not self.master.anticache
def toggle_anticomp(self):
self.master.anticomp = not self.master.anticomp
self.master.options.anticomp = not self.master.options.anticomp
def toggle_killextra(self):
self.master.killextra = not self.master.killextra

2
mitmproxy/console/statusbar.py
View File

@ -189,7 +189,7 @@ class StatusBar(urwid.WidgetWrap):
opts = []
if self.master.anticache:
opts.append("anticache")
if self.master.anticomp:
if self.master.options.anticomp:
opts.append("anticomp")
if self.master.showhost:
opts.append("showhost")

11
mitmproxy/controller.py
View File

@ -6,6 +6,8 @@ import contextlib
from six.moves import queue
from mitmproxy import addons
from mitmproxy import options
from . import ctx as mitmproxy_ctx
from netlib import basethread
from . import exceptions
@ -49,7 +51,9 @@ class Master(object):
"""
The master handles mitmproxy's main event loop.
"""
def __init__(self, *servers):
def __init__(self, opts, *servers):
self.options = opts or options.Options()
self.addons = addons.Addons(self)
self.event_queue = queue.Queue()
self.should_exit = threading.Event()
self.servers = []
@ -121,6 +125,7 @@ class Master(object):
for server in self.servers:
server.shutdown()
self.should_exit.set()
self.addons.done()
class ServerThread(basethread.BaseThread):
@ -191,6 +196,10 @@ def handler(f):
with master.handlecontext():
ret = f(master, message)
if handling:
# Python2/3 compatibility hack
fn = getattr(f, "func_name", None) or getattr(f, "__name__")
master.addons(fn, message)
if handling and not message.reply.acked and not message.reply.taken:
message.reply.ack()

24
mitmproxy/dump.py
View File

@ -12,6 +12,7 @@ from mitmproxy import exceptions
from mitmproxy import filt
from mitmproxy import flow
from mitmproxy import options
from mitmproxy import builtins
from netlib import human
from netlib import tcp
from netlib import strutils
@ -58,7 +59,8 @@ class Options(options.Options):
class DumpMaster(flow.FlowMaster):
def __init__(self, server, options, outfile=None):
flow.FlowMaster.__init__(self, server, flow.State())
flow.FlowMaster.__init__(self, options, server, flow.State())
self.addons.add(*builtins.default_addons())
self.outfile = outfile
self.o = options
self.anticache = options.anticache
@ -137,8 +139,8 @@ class DumpMaster(flow.FlowMaster):
self.add_event("Flow file corrupted.", "error")
raise DumpError(v)
if self.o.app:
self.start_app(self.o.app_host, self.o.app_port)
if self.options.app:
self.start_app(self.options.app_host, self.options.app_port)
def _readflow(self, paths):
"""
@ -152,7 +154,7 @@ class DumpMaster(flow.FlowMaster):
def add_event(self, e, level="info"):
needed = dict(error=0, info=1, debug=2).get(level, 1)
if self.o.verbosity >= needed:
if self.options.verbosity >= needed:
self.echo(
e,
fg="red" if level == "error" else None,
@ -172,7 +174,7 @@ class DumpMaster(flow.FlowMaster):
click.secho(text, file=self.outfile, **style)
def _echo_message(self, message):
if self.o.flow_detail >= 2 and hasattr(message, "headers"):
if self.options.flow_detail >= 2 and hasattr(message, "headers"):
headers = "\r\n".join(
"{}: {}".format(
click.style(strutils.bytes_to_escaped_str(k), fg="blue", bold=True),
@ -180,7 +182,7 @@ class DumpMaster(flow.FlowMaster):
for k, v in message.headers.fields
)
self.echo(headers, indent=4)
if self.o.flow_detail >= 3:
if self.options.flow_detail >= 3:
if message.content is None:
self.echo("(content missing)", indent=4)
elif message.content:
@ -213,7 +215,7 @@ class DumpMaster(flow.FlowMaster):
for (style, text) in line:
yield click.style(text, **styles.get(style, {}))
if self.o.flow_detail == 3:
if self.options.flow_detail == 3:
lines_to_echo = itertools.islice(lines, 70)
else:
lines_to_echo = lines
@ -228,7 +230,7 @@ class DumpMaster(flow.FlowMaster):
if next(lines, None):
self.echo("(cut off)", indent=4, dim=True)
if self.o.flow_detail >= 2:
if self.options.flow_detail >= 2:
self.echo("")
def _echo_request_line(self, flow):
@ -302,7 +304,7 @@ class DumpMaster(flow.FlowMaster):
self.echo(line)
def echo_flow(self, f):
if self.o.flow_detail == 0:
if self.options.flow_detail == 0:
return
if f.request:
@ -350,7 +352,7 @@ class DumpMaster(flow.FlowMaster):
def tcp_message(self, f):
super(DumpMaster, self).tcp_message(f)
if self.o.flow_detail == 0:
if self.options.flow_detail == 0:
return
message = f.messages[-1]
direction = "->" if message.from_client else "<-"
@ -362,7 +364,7 @@ class DumpMaster(flow.FlowMaster):
self._echo_message(message)
def run(self): # pragma: no cover
if self.o.rfile and not self.o.keepserving:
if self.options.rfile and not self.options.keepserving:
self.unload_scripts() # make sure to trigger script unload events.
return
super(DumpMaster, self).run()

7
mitmproxy/flow/master.py
View File

@ -27,8 +27,8 @@ class FlowMaster(controller.Master):
if len(self.servers) > 0:
return self.servers[0]
def __init__(self, server, state):
super(FlowMaster, self).__init__()
def __init__(self, options, server, state):
super(FlowMaster, self).__init__(options)
if server:
self.add_server(server)
self.state = state
@ -46,7 +46,6 @@ class FlowMaster(controller.Master):
self.stickyauth_txt = None
self.anticache = False
self.anticomp = False
self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies]
self.refresh_server_playback = False
self.replacehooks = modules.ReplaceHooks()
@ -332,8 +331,6 @@ class FlowMaster(controller.Master):
if self.anticache:
f.request.anticache()
if self.anticomp:
f.request.anticomp()
if self.server_playback:
pb = self.do_server_playback(f)

9
mitmproxy/web/master.py
View File

@ -6,6 +6,7 @@ import collections
import tornado.httpserver
import tornado.ioloop
from mitmproxy import builtins
from mitmproxy import controller
from mitmproxy import exceptions
from mitmproxy import flow
@ -147,9 +148,11 @@ class Options(options.Options):
class WebMaster(flow.FlowMaster):
def __init__(self, server, options):
self.options = options
super(WebMaster, self).__init__(server, WebState())
self.app = app.Application(self, self.options.wdebug, self.options.wauthenticator)
super(WebMaster, self).__init__(options, server, WebState())
self.addons.add(*builtins.default_addons())
self.app = app.Application(
self, self.options.wdebug, self.options.wauthenticator
)
if options.rfile:
try:
self.load_flows_file(options.rfile)

0
test/mitmproxy/builtins/__init__.py Normal file
View File

22
test/mitmproxy/builtins/test_anticomp.py Normal file
View File

@ -0,0 +1,22 @@
from .. import tutils, mastertest
from mitmproxy.builtins import anticomp
from mitmproxy.flow import master
from mitmproxy.flow import state
from mitmproxy import options
class TestAntiComp(mastertest.MasterTest):
def test_simple(self):
s = state.State()
m = master.FlowMaster(options.Options(anticomp = True), None, s)
sa = anticomp.AntiComp()
m.addons.add(sa)
f = tutils.tflow(resp=True)
self.invoke(m, "request", f)
f = tutils.tflow(resp=True)
f.request.headers["Accept-Encoding"] = "foobar"
self.invoke(m, "request", f)
assert f.request.headers["Accept-Encoding"] == "identity"

8
test/mitmproxy/mastertest.py
View File

@ -3,10 +3,16 @@ import mock
from . import tutils
import netlib.tutils
from mitmproxy import flow, proxy, models
from mitmproxy import flow, proxy, models, controller
class MasterTest:
def invoke(self, master, handler, message):
with master.handlecontext():
func = getattr(master, handler)
func(message)
message.reply = controller.DummyReply()
def cycle(self, master, content):
f = tutils.tflow(req=netlib.tutils.treq(content=content))
l = proxy.Log("connect")

20
test/mitmproxy/test_addons.py Normal file
View File

@ -0,0 +1,20 @@
from __future__ import absolute_import, print_function, division
from mitmproxy import addons
from mitmproxy import controller
from mitmproxy import options
class TAddon:
def __init__(self, name):
self.name = name
def __repr__(self):
return "Addon(%s)" % self.name
def test_simple():
m = controller.Master(options.Options())
a = addons.Addons(m)
a.add(TAddon("one"))
assert a.has_addon("one")
assert not a.has_addon("two")

4
test/mitmproxy/test_controller.py
View File

@ -25,7 +25,7 @@ class TestMaster(object):
# Speed up test
super(DummyMaster, self).tick(0)
m = DummyMaster()
m = DummyMaster(None)
assert not m.should_exit.is_set()
msg = TMsg()
msg.reply = controller.DummyReply()
@ -34,7 +34,7 @@ class TestMaster(object):
assert m.should_exit.is_set()
def test_server_simple(self):
m = controller.Master()
m = controller.Master(None)
s = DummyServer(None)
m.add_server(s)
m.start()

39
test/mitmproxy/test_flow.py
View File

@ -139,7 +139,7 @@ class TestClientPlaybackState:
def test_tick(self):
first = tutils.tflow()
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.start_client_playback([first, tutils.tflow()], True)
c = fm.client_playback
c.testing = True
@ -470,7 +470,7 @@ class TestFlow(object):
def test_kill(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow()
f.intercept(mock.Mock())
f.kill(fm)
@ -479,7 +479,7 @@ class TestFlow(object):
def test_killall(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow()
f.intercept(fm)
@ -714,7 +714,7 @@ class TestSerialize:
def test_load_flows(self):
r = self._treader()
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.load_flows(r)
assert len(s.flows) == 6
@ -725,7 +725,7 @@ class TestSerialize:
mode="reverse",
upstream_server=("https", ("use-this-domain", 80))
)
fm = flow.FlowMaster(DummyServer(conf), s)
fm = flow.FlowMaster(None, DummyServer(conf), s)
fm.load_flows(r)
assert s.flows[0].request.host == "use-this-domain"
@ -772,7 +772,7 @@ class TestFlowMaster:
def test_load_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.load_script(tutils.test_data.path("data/scripts/a.py"))
fm.load_script(tutils.test_data.path("data/scripts/a.py"))
@ -788,14 +788,14 @@ class TestFlowMaster:
def test_getset_ignore(self):
p = mock.Mock()
p.config.check_ignore = HostMatcher()
fm = flow.FlowMaster(p, flow.State())
fm = flow.FlowMaster(None, p, flow.State())
assert not fm.get_ignore_filter()
fm.set_ignore_filter(["^apple\.com:", ":443$"])
assert fm.get_ignore_filter()
def test_replay(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow(resp=True)
f.request.content = None
assert "missing" in fm.replay_request(f)
@ -808,7 +808,7 @@ class TestFlowMaster:
def test_script_reqerr(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.load_script(tutils.test_data.path("data/scripts/reqerr.py"))
f = tutils.tflow()
fm.clientconnect(f.client_conn)
@ -816,7 +816,7 @@ class TestFlowMaster:
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.load_script(tutils.test_data.path("data/scripts/all.py"))
f = tutils.tflow(resp=True)
@ -852,7 +852,7 @@ class TestFlowMaster:
def test_duplicate_flow(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow(resp=True)
fm.load_flow(f)
assert s.flow_count() == 1
@ -863,14 +863,13 @@ class TestFlowMaster:
def test_create_flow(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
assert fm.create_request("GET", "http", "example.com", 80, "/")
def test_all(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.anticache = True
fm.anticomp = True
f = tutils.tflow(req=None)
fm.clientconnect(f.client_conn)
f.request = HTTPRequest.wrap(netlib.tutils.treq())
@ -895,7 +894,7 @@ class TestFlowMaster:
f = tutils.tflow(resp=True)
pb = [tutils.tflow(resp=True), f]
fm = flow.FlowMaster(DummyServer(ProxyConfig()), s)
fm = flow.FlowMaster(None, DummyServer(ProxyConfig()), s)
assert not fm.start_server_playback(
pb,
False,
@ -923,7 +922,7 @@ class TestFlowMaster:
f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request))
pb = [f]
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.refresh_server_playback = True
assert not fm.do_server_playback(tutils.tflow())
@ -965,7 +964,7 @@ class TestFlowMaster:
f = tutils.tflow()
f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request))
pb = [f]
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.refresh_server_playback = True
fm.start_server_playback(
pb,
@ -985,7 +984,7 @@ class TestFlowMaster:
def test_stickycookie(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
assert "Invalid" in fm.set_stickycookie("~h")
fm.set_stickycookie(".*")
assert fm.stickycookie_state
@ -1007,7 +1006,7 @@ class TestFlowMaster:
def test_stickyauth(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
assert "Invalid" in fm.set_stickyauth("~h")
fm.set_stickyauth(".*")
assert fm.stickyauth_state
@ -1035,7 +1034,7 @@ class TestFlowMaster:
return list(r.stream())
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
f = tutils.tflow(resp=True)
with open(p, "ab") as tmpfile:

2
test/mitmproxy/test_script.py
View File

@ -4,7 +4,7 @@ from . import tutils
def test_duplicate_flow():
s = flow.State()
fm = flow.FlowMaster(None, s)
fm = flow.FlowMaster(None, None, s)
fm.load_script(tutils.test_data.path("data/scripts/duplicate_flow.py"))
f = tutils.tflow()
fm.request(f)

2
test/mitmproxy/tservers.py
View File

@ -34,7 +34,7 @@ class TestMaster(flow.FlowMaster):
config.port = 0
s = ProxyServer(config)
state = flow.State()
flow.FlowMaster.__init__(self, s, state)
flow.FlowMaster.__init__(self, None, s, state)
self.apps.add(testapp, "testapp", 80)
self.apps.add(errapp, "errapp", 80)
self.clear_log()