mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
Basic outline of addons
Add addons.py, integrate with our event mechanism, and change the Master API so options is the first init argument.
This commit is contained in:
Aldo Cortesi
parent
c9a0fe6a0e
commit
7f8fd3cdff
63
mitmproxy/addons.py
Normal file
63
mitmproxy/addons.py
Normal 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)
|
||||
@ -216,10 +216,9 @@ class ConsoleMaster(flow.FlowMaster):
|
||||
palette = []
|
||||
|
||||
def __init__(self, server, options):
|
||||
flow.FlowMaster.__init__(self, server, ConsoleState())
|
||||
flow.FlowMaster.__init__(self, options, server, ConsoleState())
|
||||
|
||||
self.stream_path = None
|
||||
self.options = options
|
||||
self.options.errored.connect(self.options_error)
|
||||
|
||||
if options.replacements:
|
||||
|
||||
@ -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)
|
||||
|
||||
if handling and not message.reply.acked and not message.reply.taken:
|
||||
message.reply.ack()
|
||||
|
||||
@ -58,7 +58,7 @@ 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.outfile = outfile
|
||||
self.o = options
|
||||
self.anticache = options.anticache
|
||||
@ -137,8 +137,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 +152,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 +172,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 +180,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 +213,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 +228,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 +302,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 +350,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 +362,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()
|
||||
|
||||
@ -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
|
||||
|
||||
@ -42,7 +42,10 @@ class Options(object):
|
||||
return self.__class__(**self._opts)
|
||||
|
||||
def __getattr__(self, attr):
|
||||
return self._opts[attr]
|
||||
if attr in self._opts:
|
||||
return self._opts[attr]
|
||||
else:
|
||||
raise AttributeError()
|
||||
|
||||
def __setattr__(self, attr, value):
|
||||
if attr not in self._opts:
|
||||
|
||||
@ -147,9 +147,10 @@ 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.app = app.Application(
|
||||
self, self.options.wdebug, self.options.wauthenticator
|
||||
)
|
||||
if options.rfile:
|
||||
try:
|
||||
self.load_flows_file(options.rfile)
|
||||
|
||||
20
test/mitmproxy/test_addons.py
Normal file
20
test/mitmproxy/test_addons.py
Normal 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")
|
||||
@ -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()
|
||||
|
||||
@ -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,12 +863,12 @@ 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)
|
||||
@ -895,7 +895,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 +923,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 +965,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 +985,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 +1007,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 +1035,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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user