mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2025-01-30 14:58:38 +00:00
Remove test handlers by using taddons.RecordingMaster
This also means expanding and tweaking the recording master API, which we reflect through the current test suite
This commit is contained in:
parent
85ddc5056b
commit
1410cbb4b6
@ -1,43 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
This example shows how to build a proxy based on mitmproxy's Flow
|
||||
primitives.
|
||||
|
||||
Heads Up: In the majority of cases, you want to use inline scripts.
|
||||
|
||||
Note that request and response messages are not automatically replied to,
|
||||
so we need to implement handlers to do this.
|
||||
"""
|
||||
from mitmproxy import controller, options, master
|
||||
from mitmproxy.proxy import ProxyServer, ProxyConfig
|
||||
|
||||
|
||||
class MyMaster(master.Master):
|
||||
def run(self):
|
||||
try:
|
||||
master.Master.run(self)
|
||||
except KeyboardInterrupt:
|
||||
self.shutdown()
|
||||
|
||||
@controller.handler
|
||||
def request(self, f):
|
||||
print("request", f)
|
||||
|
||||
@controller.handler
|
||||
def response(self, f):
|
||||
print("response", f)
|
||||
|
||||
@controller.handler
|
||||
def error(self, f):
|
||||
print("error", f)
|
||||
|
||||
@controller.handler
|
||||
def log(self, l):
|
||||
print("log", l.msg)
|
||||
|
||||
|
||||
opts = options.Options(cadir="~/.mitmproxy/")
|
||||
config = ProxyConfig(opts)
|
||||
server = ProxyServer(config)
|
||||
m = MyMaster(opts, server)
|
||||
m.run()
|
@ -1,5 +1,6 @@
|
||||
import mitmproxy
|
||||
from mitmproxy.net import tcp
|
||||
from mitmproxy import ctx
|
||||
|
||||
|
||||
class CheckALPN:
|
||||
@ -9,9 +10,8 @@ class CheckALPN:
|
||||
def configure(self, options, updated):
|
||||
self.failed = mitmproxy.ctx.master.options.http2 and not tcp.HAS_ALPN
|
||||
if self.failed:
|
||||
mitmproxy.ctx.master.add_log(
|
||||
ctx.log.warn(
|
||||
"HTTP/2 is disabled because ALPN support missing!\n"
|
||||
"OpenSSL 1.0.2+ required to support HTTP/2 connections.\n"
|
||||
"Use --no-http2 to silence this warning.",
|
||||
"warn",
|
||||
"Use --no-http2 to silence this warning."
|
||||
)
|
||||
|
@ -12,7 +12,10 @@ class _AddonWrapper:
|
||||
self.addons = addons
|
||||
|
||||
def trigger(self, event, *args, **kwargs):
|
||||
self.master.events.append((event, args, kwargs))
|
||||
if event == "log":
|
||||
self.master.logs.append(args[0])
|
||||
else:
|
||||
self.master.events.append((event, args, kwargs))
|
||||
return self.addons.trigger(event, *args, **kwargs)
|
||||
|
||||
def __getattr__(self, attr):
|
||||
@ -26,13 +29,19 @@ class RecordingMaster(mitmproxy.master.Master):
|
||||
self.events = []
|
||||
self.logs = []
|
||||
|
||||
def has_log(self, txt, level=None):
|
||||
for i in self.logs:
|
||||
if level and i.level != level:
|
||||
continue
|
||||
if txt.lower() in i.msg.lower():
|
||||
return True
|
||||
return False
|
||||
|
||||
def has_event(self, name):
|
||||
for i in self.events:
|
||||
if i[0] == name:
|
||||
return True
|
||||
|
||||
def add_log(self, e, level):
|
||||
self.logs.append((level, e))
|
||||
return False
|
||||
|
||||
def clear(self):
|
||||
self.logs = []
|
||||
|
@ -13,7 +13,6 @@ import traceback
|
||||
import urwid
|
||||
|
||||
from mitmproxy import addons
|
||||
from mitmproxy import controller
|
||||
from mitmproxy import exceptions
|
||||
from mitmproxy import master
|
||||
from mitmproxy import io
|
||||
|
@ -1,4 +1,3 @@
|
||||
from mitmproxy import ctx
|
||||
from mitmproxy import addons
|
||||
from mitmproxy import options
|
||||
from mitmproxy import master
|
||||
|
@ -12,7 +12,7 @@ class TestCheckALPN:
|
||||
with taddons.context() as tctx:
|
||||
a = check_alpn.CheckALPN()
|
||||
tctx.configure(a)
|
||||
assert not any(msg in m for l, m in tctx.master.logs)
|
||||
assert not tctx.master.has_log(msg)
|
||||
|
||||
def test_check_no_alpn(self, disable_alpn):
|
||||
msg = 'ALPN support missing'
|
||||
@ -20,4 +20,4 @@ class TestCheckALPN:
|
||||
with taddons.context() as tctx:
|
||||
a = check_alpn.CheckALPN()
|
||||
tctx.configure(a)
|
||||
assert any(msg in m for l, m in tctx.master.logs)
|
||||
assert tctx.master.has_log(msg)
|
||||
|
@ -16,4 +16,4 @@ class TestCheckCA:
|
||||
tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock(return_value=expired)
|
||||
a = check_ca.CheckCA()
|
||||
tctx.configure(a)
|
||||
assert any(msg in m for l, m in tctx.master.logs) is expired
|
||||
assert tctx.master.has_log(msg) is expired
|
||||
|
@ -151,7 +151,7 @@ class TestContentView:
|
||||
with taddons.context(options=options.Options()) as ctx:
|
||||
ctx.configure(d, flow_detail=4, verbosity=3)
|
||||
d.response(tflow.tflow())
|
||||
assert "Content viewer failed" in ctx.master.logs[0][1]
|
||||
assert ctx.master.has_log("content viewer failed")
|
||||
|
||||
|
||||
def test_tcp():
|
||||
|
@ -22,14 +22,12 @@ def test_scriptenv():
|
||||
with taddons.context() as tctx:
|
||||
with script.scriptenv("path", []):
|
||||
raise SystemExit
|
||||
assert tctx.master.logs[0][0] == "error"
|
||||
assert "exited" in tctx.master.logs[0][1]
|
||||
assert tctx.master.has_log("exited", "error")
|
||||
|
||||
tctx.master.clear()
|
||||
with script.scriptenv("path", []):
|
||||
raise ValueError("fooo")
|
||||
assert tctx.master.logs[0][0] == "error"
|
||||
assert "foo" in tctx.master.logs[0][1]
|
||||
assert tctx.master.has_log("fooo", "error")
|
||||
|
||||
|
||||
class Called:
|
||||
@ -147,11 +145,11 @@ class TestScript:
|
||||
sc.start(tctx.options)
|
||||
f = tflow.tflow(resp=True)
|
||||
sc.request(f)
|
||||
assert tctx.master.logs[0][0] == "error"
|
||||
assert len(tctx.master.logs[0][1].splitlines()) == 6
|
||||
assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0][1])
|
||||
assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0][1])
|
||||
assert tctx.master.logs[0][1].endswith("ValueError: Error!\n")
|
||||
assert tctx.master.logs[0].level == "error"
|
||||
assert len(tctx.master.logs[0].msg.splitlines()) == 6
|
||||
assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0].msg)
|
||||
assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0].msg)
|
||||
assert tctx.master.logs[0].msg.endswith("ValueError: Error!\n")
|
||||
|
||||
def test_addon(self):
|
||||
with taddons.context() as tctx:
|
||||
@ -256,19 +254,19 @@ class TestScriptLoader:
|
||||
"%s %s" % (rec, "c"),
|
||||
]
|
||||
)
|
||||
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
|
||||
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
|
||||
assert debug == [
|
||||
('debug', 'a start'),
|
||||
('debug', 'a configure'),
|
||||
('debug', 'a running'),
|
||||
'a start',
|
||||
'a configure',
|
||||
'a running',
|
||||
|
||||
('debug', 'b start'),
|
||||
('debug', 'b configure'),
|
||||
('debug', 'b running'),
|
||||
'b start',
|
||||
'b configure',
|
||||
'b running',
|
||||
|
||||
('debug', 'c start'),
|
||||
('debug', 'c configure'),
|
||||
('debug', 'c running'),
|
||||
'c start',
|
||||
'c configure',
|
||||
'c running',
|
||||
]
|
||||
tctx.master.logs = []
|
||||
tctx.configure(
|
||||
@ -279,8 +277,7 @@ class TestScriptLoader:
|
||||
"%s %s" % (rec, "b"),
|
||||
]
|
||||
)
|
||||
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
|
||||
# No events, only order has changed
|
||||
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
|
||||
assert debug == []
|
||||
|
||||
tctx.master.logs = []
|
||||
@ -291,11 +288,11 @@ class TestScriptLoader:
|
||||
"%s %s" % (rec, "a"),
|
||||
]
|
||||
)
|
||||
debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
|
||||
debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
|
||||
assert debug == [
|
||||
('debug', 'c done'),
|
||||
('debug', 'b done'),
|
||||
('debug', 'x start'),
|
||||
('debug', 'x configure'),
|
||||
('debug', 'x running'),
|
||||
'c done',
|
||||
'b done',
|
||||
'x start',
|
||||
'x configure',
|
||||
'x running',
|
||||
]
|
||||
|
@ -250,17 +250,12 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
|
||||
assert p.request(req)
|
||||
|
||||
def test_get_connection_switching(self):
|
||||
def switched(l):
|
||||
for i in l:
|
||||
if "serverdisconnect" in i:
|
||||
return True
|
||||
|
||||
req = "get:'%s/p/200:b@1'"
|
||||
p = self.pathoc()
|
||||
with p.connect():
|
||||
assert p.request(req % self.server.urlbase)
|
||||
assert p.request(req % self.server2.urlbase)
|
||||
assert switched(self.proxy.tlog)
|
||||
assert self.proxy.tmaster.has_log("serverdisconnect")
|
||||
|
||||
def test_blank_leading_line(self):
|
||||
p = self.pathoc()
|
||||
@ -602,7 +597,7 @@ class TestHttps2Http(tservers.ReverseProxyTest):
|
||||
p = self.pathoc(ssl=True, sni="example.com")
|
||||
with p.connect():
|
||||
assert p.request("get:'/p/200'").status_code == 200
|
||||
assert all("Error in handle_sni" not in msg for msg in self.proxy.tlog)
|
||||
assert not self.proxy.tmaster.has_log("error in handle_sni")
|
||||
|
||||
def test_http(self):
|
||||
p = self.pathoc(ssl=False)
|
||||
@ -877,8 +872,7 @@ class TestServerConnect(tservers.HTTPProxyTest):
|
||||
def test_unnecessary_serverconnect(self):
|
||||
"""A replayed/fake response with no upstream_cert should not connect to an upstream server"""
|
||||
assert self.pathod("200").status_code == 200
|
||||
for msg in self.proxy.tmaster.tlog:
|
||||
assert "serverconnect" not in msg
|
||||
assert not self.proxy.tmaster.has_log("serverconnect")
|
||||
|
||||
|
||||
class MasterKillRequest(tservers.TestMaster):
|
||||
|
@ -43,7 +43,7 @@ class TestConcurrent(tservers.MasterTest):
|
||||
)
|
||||
)
|
||||
sc.start(tctx.options)
|
||||
assert "decorator not supported" in tctx.master.logs[0][1]
|
||||
assert tctx.master.has_log("decorator not supported")
|
||||
|
||||
def test_concurrent_class(self):
|
||||
with taddons.context() as tctx:
|
||||
|
@ -7,6 +7,7 @@ from mitmproxy.exceptions import Kill, ControlException
|
||||
from mitmproxy import controller
|
||||
from mitmproxy import master
|
||||
from mitmproxy import proxy
|
||||
from mitmproxy.test import taddons
|
||||
|
||||
|
||||
class TMsg:
|
||||
@ -15,22 +16,18 @@ class TMsg:
|
||||
|
||||
class TestMaster:
|
||||
def test_simple(self):
|
||||
class DummyMaster(master.Master):
|
||||
@controller.handler
|
||||
class tAddon:
|
||||
def log(self, _):
|
||||
m.should_exit.set()
|
||||
ctx.master.should_exit.set()
|
||||
|
||||
def tick(self, timeout):
|
||||
# Speed up test
|
||||
super().tick(0)
|
||||
|
||||
m = DummyMaster(None, proxy.DummyServer(None))
|
||||
assert not m.should_exit.is_set()
|
||||
msg = TMsg()
|
||||
msg.reply = controller.DummyReply()
|
||||
m.event_queue.put(("log", msg))
|
||||
m.run()
|
||||
assert m.should_exit.is_set()
|
||||
with taddons.context() as ctx:
|
||||
ctx.master.addons.add(tAddon())
|
||||
assert not ctx.master.should_exit.is_set()
|
||||
msg = TMsg()
|
||||
msg.reply = controller.DummyReply()
|
||||
ctx.master.event_queue.put(("log", msg))
|
||||
ctx.master.run()
|
||||
assert ctx.master.should_exit.is_set()
|
||||
|
||||
def test_server_simple(self):
|
||||
m = master.Master(None, proxy.DummyServer(None))
|
||||
|
@ -6,7 +6,6 @@ import sys
|
||||
import mitmproxy.platform
|
||||
from mitmproxy.proxy.config import ProxyConfig
|
||||
from mitmproxy.proxy.server import ProxyServer
|
||||
from mitmproxy import master
|
||||
from mitmproxy import controller
|
||||
from mitmproxy import options
|
||||
from mitmproxy import exceptions
|
||||
@ -17,6 +16,7 @@ import pathod.pathoc
|
||||
|
||||
from mitmproxy.test import tflow
|
||||
from mitmproxy.test import tutils
|
||||
from mitmproxy.test import taddons
|
||||
|
||||
|
||||
class MasterTest:
|
||||
@ -68,11 +68,11 @@ class TestState:
|
||||
# self.flows.append(f)
|
||||
|
||||
|
||||
class TestMaster(master.Master):
|
||||
class TestMaster(taddons.RecordingMaster):
|
||||
|
||||
def __init__(self, opts, config):
|
||||
s = ProxyServer(config)
|
||||
master.Master.__init__(self, opts, s)
|
||||
super().__init__(opts, s)
|
||||
|
||||
def clear_addons(self, addons):
|
||||
self.addons.clear()
|
||||
@ -82,16 +82,9 @@ class TestMaster(master.Master):
|
||||
self.addons.configure_all(self.options, self.options.keys())
|
||||
self.addons.trigger("running")
|
||||
|
||||
def clear_log(self):
|
||||
self.tlog = []
|
||||
|
||||
def reset(self, addons):
|
||||
self.clear_addons(addons)
|
||||
self.clear_log()
|
||||
|
||||
@controller.handler
|
||||
def log(self, e):
|
||||
self.tlog.append(e.msg)
|
||||
self.clear()
|
||||
|
||||
|
||||
class ProxyThread(threading.Thread):
|
||||
@ -111,7 +104,7 @@ class ProxyThread(threading.Thread):
|
||||
|
||||
@property
|
||||
def tlog(self):
|
||||
return self.tmaster.tlog
|
||||
return self.tmaster.logs
|
||||
|
||||
def run(self):
|
||||
self.tmaster.run()
|
||||
|
Loading…
Reference in New Issue
Block a user