Remove watchdog, solidify script testing

- Remove the watchdog dependency. We now just stat the script file every 2
seconds to check for an updated mtime.
- Further solidify our script testing, and in particular make the example tests
nicer. These should exemplify how we want users to test their own addon
scripts. More work on addon testing to follow.
This commit is contained in:
Aldo Cortesi 2017-04-26 07:13:36 +12:00 committed by Aldo Cortesi
parent e6eeab6094
commit 7aa2081894
9 changed files with 118 additions and 184 deletions

View File

@ -129,9 +129,13 @@ class AddonManager:
def register(self, addon): def register(self, addon):
""" """
Register an addon and all its sub-addons with the manager without Register an addon, call its load event, and then register all its
adding it to the chain. This should be used by addons that sub-addons. This should be used by addons that dynamically manage
dynamically manage addons. Must be called within a current context. addons.
If the calling addon is already running, it should follow with
running and configure events. Must be called within a current
context.
""" """
for a in traverse([addon]): for a in traverse([addon]):
name = _get_name(a) name = _get_name(a)

View File

@ -1,15 +1,12 @@
import os import os
import importlib import importlib
import threading import time
import sys import sys
from mitmproxy import addonmanager from mitmproxy import addonmanager
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy import ctx from mitmproxy import ctx
import watchdog.events
from watchdog.observers import polling
def load_script(actx, path): def load_script(actx, path):
if not os.path.exists(path): if not os.path.exists(path):
@ -28,60 +25,30 @@ def load_script(actx, path):
sys.path[:] = oldpath sys.path[:] = oldpath
class ReloadHandler(watchdog.events.FileSystemEventHandler):
def __init__(self, callback):
self.callback = callback
def filter(self, event):
"""
Returns True only when .py file is changed
"""
if event.is_directory:
return False
if os.path.basename(event.src_path).startswith("."):
return False
if event.src_path.endswith(".py"):
return True
return False
def on_modified(self, event):
if self.filter(event):
self.callback()
def on_created(self, event):
if self.filter(event):
self.callback()
class Script: class Script:
""" """
An addon that manages a single script. An addon that manages a single script.
""" """
ReloadInterval = 2
def __init__(self, path): def __init__(self, path):
self.name = "scriptmanager:" + path self.name = "scriptmanager:" + path
self.path = path self.path = path
self.ns = None self.ns = None
self.observer = None
self.last_options = None self.last_options = None
self.should_reload = threading.Event() self.last_load = 0
self.last_mtime = 0
def load(self, l):
self.ns = load_script(ctx, self.path)
@property @property
def addons(self): def addons(self):
if self.ns is not None: return [self.ns] if self.ns else []
return [self.ns]
return []
def reload(self):
self.should_reload.set()
def tick(self): def tick(self):
if self.should_reload.is_set(): if time.time() - self.last_load > self.ReloadInterval:
self.should_reload.clear() mtime = os.stat(self.path).st_mtime
ctx.log.info("Reloading script: %s" % self.name) if mtime > self.last_mtime:
ctx.log.info("Loading script: %s" % self.name)
if self.ns: if self.ns:
ctx.master.addons.remove(self.ns) ctx.master.addons.remove(self.ns)
self.ns = load_script(ctx, self.path) self.ns = load_script(ctx, self.path)
@ -89,18 +56,18 @@ class Script:
# We're already running, so we have to explicitly register and # We're already running, so we have to explicitly register and
# configure the addon # configure the addon
ctx.master.addons.register(self.ns) ctx.master.addons.register(self.ns)
self.configure(self.last_options, self.last_options.keys()) ctx.master.addons.invoke_addon(self.ns, "running")
ctx.master.addons.invoke_addon(
self.ns,
"configure",
self.last_options,
self.last_options.keys()
)
self.last_load = time.time()
self.last_mtime = mtime
def configure(self, options, updated): def configure(self, options, updated):
self.last_options = options self.last_options = options
if not self.observer:
self.observer = polling.PollingObserver()
# Bind the handler to the real underlying master object
self.observer.schedule(
ReloadHandler(self.reload),
os.path.dirname(self.path) or "."
)
self.observer.start()
class ScriptLoader: class ScriptLoader:

View File

@ -112,9 +112,12 @@ class context:
) )
def script(self, path): def script(self, path):
"""
Loads a script from path, and returns the enclosed addon.
"""
sc = script.Script(path) sc = script.Script(path)
loader = addonmanager.Loader(self.master) loader = addonmanager.Loader(self.master)
sc.load(loader) self.master.addons.invoke_addon(sc, "load", loader)
for a in addonmanager.traverse(sc.addons): self.configure(sc)
getattr(a, "load", lambda x: None)(loader) self.master.addons.invoke_addon(sc, "tick")
return sc return sc.addons[0] if sc.addons else None

View File

@ -80,7 +80,6 @@ setup(
"ruamel.yaml>=0.13.2, <0.15", "ruamel.yaml>=0.13.2, <0.15",
"tornado>=4.3, <4.6", "tornado>=4.3, <4.6",
"urwid>=1.3.1, <1.4", "urwid>=1.3.1, <1.4",
"watchdog>=0.8.3, <0.9",
"brotlipy>=0.5.1, <0.7", "brotlipy>=0.5.1, <0.7",
"sortedcontainers>=1.5.4, <1.6", "sortedcontainers>=1.5.4, <1.6",
# transitive from cryptography, we just blacklist here. # transitive from cryptography, we just blacklist here.

View File

@ -1,9 +1,4 @@
from mitmproxy import options
from mitmproxy import contentviews from mitmproxy import contentviews
from mitmproxy import proxy
from mitmproxy import master
from mitmproxy.addons import script
from mitmproxy.test import tflow from mitmproxy.test import tflow
from mitmproxy.test import tutils from mitmproxy.test import tutils
from mitmproxy.test import taddons from mitmproxy.test import taddons
@ -14,34 +9,17 @@ from ..mitmproxy import tservers
example_dir = tutils.test_data.push("../examples") example_dir = tutils.test_data.push("../examples")
class ScriptError(Exception):
pass
class RaiseMaster(master.Master):
def add_log(self, e, level):
if level in ("warn", "error"):
raise ScriptError(e)
def tscript(cmd, args=""):
o = options.Options()
cmd = example_dir.path(cmd)
m = RaiseMaster(o, proxy.DummyServer())
sc = script.Script(cmd)
m.addons.add(sc)
return m, sc
class TestScripts(tservers.MasterTest): class TestScripts(tservers.MasterTest):
def test_add_header(self): def test_add_header(self):
m, _ = tscript("simple/add_header.py") with taddons.context() as tctx:
a = tctx.script(example_dir.path("simple/add_header.py"))
f = tflow.tflow(resp=tutils.tresp()) f = tflow.tflow(resp=tutils.tresp())
m.addons.handle_lifecycle("response", f) a.response(f)
assert f.response.headers["newheader"] == "foo" assert f.response.headers["newheader"] == "foo"
def test_custom_contentviews(self): def test_custom_contentviews(self):
m, sc = tscript("simple/custom_contentview.py") with taddons.context() as tctx:
tctx.script(example_dir.path("simple/custom_contentview.py"))
swapcase = contentviews.get("swapcase") swapcase = contentviews.get("swapcase")
_, fmt = swapcase(b"<html>Test!</html>") _, fmt = swapcase(b"<html>Test!</html>")
assert any(b'tEST!' in val[0][1] for val in fmt) assert any(b'tEST!' in val[0][1] for val in fmt)
@ -61,55 +39,61 @@ class TestScripts(tservers.MasterTest):
assert b'iframe' in content and b'evil_iframe' in content assert b'iframe' in content and b'evil_iframe' in content
def test_modify_form(self): def test_modify_form(self):
m, sc = tscript("simple/modify_form.py") with taddons.context() as tctx:
sc = tctx.script(example_dir.path("simple/modify_form.py"))
form_header = Headers(content_type="application/x-www-form-urlencoded") form_header = Headers(content_type="application/x-www-form-urlencoded")
f = tflow.tflow(req=tutils.treq(headers=form_header)) f = tflow.tflow(req=tutils.treq(headers=form_header))
m.addons.handle_lifecycle("request", f) sc.request(f)
assert f.request.urlencoded_form["mitmproxy"] == "rocks" assert f.request.urlencoded_form["mitmproxy"] == "rocks"
f.request.headers["content-type"] = "" f.request.headers["content-type"] = ""
m.addons.handle_lifecycle("request", f) sc.request(f)
assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] assert list(f.request.urlencoded_form.items()) == [("foo", "bar")]
def test_modify_querystring(self): def test_modify_querystring(self):
m, sc = tscript("simple/modify_querystring.py") with taddons.context() as tctx:
sc = tctx.script(example_dir.path("simple/modify_querystring.py"))
f = tflow.tflow(req=tutils.treq(path="/search?q=term")) f = tflow.tflow(req=tutils.treq(path="/search?q=term"))
m.addons.handle_lifecycle("request", f) sc.request(f)
assert f.request.query["mitmproxy"] == "rocks" assert f.request.query["mitmproxy"] == "rocks"
f.request.path = "/" f.request.path = "/"
m.addons.handle_lifecycle("request", f) sc.request(f)
assert f.request.query["mitmproxy"] == "rocks" assert f.request.query["mitmproxy"] == "rocks"
def test_redirect_requests(self): def test_redirect_requests(self):
m, sc = tscript("simple/redirect_requests.py") with taddons.context() as tctx:
sc = tctx.script(example_dir.path("simple/redirect_requests.py"))
f = tflow.tflow(req=tutils.treq(host="example.org")) f = tflow.tflow(req=tutils.treq(host="example.org"))
m.addons.handle_lifecycle("request", f) sc.request(f)
assert f.request.host == "mitmproxy.org" assert f.request.host == "mitmproxy.org"
def test_send_reply_from_proxy(self): def test_send_reply_from_proxy(self):
m, sc = tscript("simple/send_reply_from_proxy.py") with taddons.context() as tctx:
sc = tctx.script(example_dir.path("simple/send_reply_from_proxy.py"))
f = tflow.tflow(req=tutils.treq(host="example.com", port=80)) f = tflow.tflow(req=tutils.treq(host="example.com", port=80))
m.addons.handle_lifecycle("request", f) sc.request(f)
assert f.response.content == b"Hello World" assert f.response.content == b"Hello World"
def test_dns_spoofing(self): def test_dns_spoofing(self):
m, sc = tscript("complex/dns_spoofing.py") with taddons.context() as tctx:
sc = tctx.script(example_dir.path("complex/dns_spoofing.py"))
original_host = "example.com" original_host = "example.com"
host_header = Headers(host=original_host) host_header = Headers(host=original_host)
f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) f = tflow.tflow(req=tutils.treq(headers=host_header, port=80))
m.addons.handle_lifecycle("requestheaders", f) tctx.master.addons.invoke_addon(sc, "requestheaders", f)
# Rewrite by reverse proxy mode # Rewrite by reverse proxy mode
f.request.scheme = "https" f.request.scheme = "https"
f.request.port = 443 f.request.port = 443
m.addons.handle_lifecycle("request", f) tctx.master.addons.invoke_addon(sc, "request", f)
assert f.request.scheme == "http" assert f.request.scheme == "http"
assert f.request.port == 80 assert f.request.port == 80

View File

@ -1,7 +1,6 @@
import traceback import traceback
import sys import sys
import time import time
import watchdog.events
import pytest import pytest
from unittest import mock from unittest import mock
@ -16,34 +15,6 @@ from mitmproxy import master
from mitmproxy.addons import script from mitmproxy.addons import script
class Called:
def __init__(self):
self.called = False
def __call__(self, *args, **kwargs):
self.called = True
def test_reloadhandler():
rh = script.ReloadHandler(Called())
assert not rh.filter(watchdog.events.DirCreatedEvent("path"))
assert not rh.filter(watchdog.events.FileModifiedEvent("/foo/.bar"))
assert not rh.filter(watchdog.events.FileModifiedEvent("/foo/bar"))
assert rh.filter(watchdog.events.FileModifiedEvent("/foo/bar.py"))
assert not rh.callback.called
rh.on_modified(watchdog.events.FileModifiedEvent("/foo/bar"))
assert not rh.callback.called
rh.on_modified(watchdog.events.FileModifiedEvent("/foo/bar.py"))
assert rh.callback.called
rh.callback.called = False
rh.on_created(watchdog.events.FileCreatedEvent("foo"))
assert not rh.callback.called
rh.on_created(watchdog.events.FileCreatedEvent("foo.py"))
assert rh.callback.called
def test_load_script(): def test_load_script():
with taddons.context() as tctx: with taddons.context() as tctx:
ns = script.load_script( ns = script.load_script(
@ -89,6 +60,8 @@ class TestScript:
) )
) )
tctx.master.addons.add(sc) tctx.master.addons.add(sc)
tctx.configure(sc)
sc.tick()
rec = tctx.master.addons.get("recorder") rec = tctx.master.addons.get("recorder")
@ -107,10 +80,12 @@ class TestScript:
f.write("\n") f.write("\n")
sc = script.Script(str(f)) sc = script.Script(str(f))
tctx.configure(sc) tctx.configure(sc)
for _ in range(5): sc.tick()
sc.reload() for _ in range(3):
sc.last_load, sc.last_mtime = 0, 0
sc.tick() sc.tick()
time.sleep(0.1) time.sleep(0.1)
tctx.master.has_log("Loading")
def test_exception(self): def test_exception(self):
with taddons.context() as tctx: with taddons.context() as tctx:
@ -118,10 +93,12 @@ class TestScript:
tutils.test_data.path("mitmproxy/data/addonscripts/error.py") tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
) )
tctx.master.addons.add(sc) tctx.master.addons.add(sc)
tctx.configure(sc)
sc.tick()
f = tflow.tflow(resp=True) f = tflow.tflow(resp=True)
tctx.master.addons.trigger("request", f) tctx.master.addons.trigger("request", f)
assert tctx.master.logs[0].level == "error"
tctx.master.has_log("ValueError: Error!") tctx.master.has_log("ValueError: Error!")
tctx.master.has_log("error.py") tctx.master.has_log("error.py")
@ -133,8 +110,10 @@ class TestScript:
) )
) )
tctx.master.addons.add(sc) tctx.master.addons.add(sc)
tctx.configure(sc)
sc.tick()
assert sc.ns.event_log == [ assert sc.ns.event_log == [
'scriptload', 'addonload' 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure'
] ]
@ -207,21 +186,23 @@ class TestScriptLoader:
"%s/c.py" % rec, "%s/c.py" % rec,
] ]
) )
tctx.master.addons.invoke_addon(sc, "tick")
debug = [i.msg for i in tctx.master.logs if i.level == "debug"] debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == [ assert debug == [
'a load', 'a load',
'a running', 'a running',
'a configure',
'a tick',
'b load', 'b load',
'b running', 'b running',
'b configure',
'b tick',
'c load', 'c load',
'c running', 'c running',
'a configure',
'b configure',
'c configure', 'c configure',
'c tick',
] ]
tctx.master.logs = [] tctx.master.logs = []
@ -233,6 +214,7 @@ class TestScriptLoader:
"%s/b.py" % rec, "%s/b.py" % rec,
] ]
) )
debug = [i.msg for i in tctx.master.logs if i.level == "debug"] debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == [ assert debug == [
'c configure', 'c configure',
@ -248,13 +230,16 @@ class TestScriptLoader:
"%s/a.py" % rec, "%s/a.py" % rec,
] ]
) )
tctx.master.addons.invoke_addon(sc, "tick")
debug = [i.msg for i in tctx.master.logs if i.level == "debug"] debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
assert debug == [ assert debug == [
'c done', 'c done',
'b done', 'b done',
'a configure',
'e load', 'e load',
'e running', 'e running',
'e configure', 'e configure',
'a configure', 'e tick',
'a tick',
] ]

View File

@ -14,7 +14,7 @@ class Addon:
def configure(options, updated): def configure(options, updated):
event_log.append("addonconfigure") event_log.append("scriptconfigure")
def load(l): def load(l):

View File

@ -2,5 +2,5 @@ from mitmproxy.script import concurrent
@concurrent @concurrent
def start(opts): def load(v):
pass pass

View File

@ -2,10 +2,7 @@ from mitmproxy.test import tflow
from mitmproxy.test import tutils from mitmproxy.test import tutils
from mitmproxy.test import taddons from mitmproxy.test import taddons
from mitmproxy import addonmanager
from mitmproxy import controller from mitmproxy import controller
from mitmproxy.addons import script
import time import time
from .. import tservers from .. import tservers
@ -36,25 +33,20 @@ class TestConcurrent(tservers.MasterTest):
def test_concurrent_err(self): def test_concurrent_err(self):
with taddons.context() as tctx: with taddons.context() as tctx:
sc = script.Script( tctx.script(
tutils.test_data.path( tutils.test_data.path(
"mitmproxy/data/addonscripts/concurrent_decorator_err.py" "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
) )
) )
l = addonmanager.Loader(tctx.master)
sc.load(l)
assert tctx.master.has_log("decorator not supported") assert tctx.master.has_log("decorator not supported")
def test_concurrent_class(self): def test_concurrent_class(self):
with taddons.context() as tctx: with taddons.context() as tctx:
sc = script.Script( sc = tctx.script(
tutils.test_data.path( tutils.test_data.path(
"mitmproxy/data/addonscripts/concurrent_decorator_class.py" "mitmproxy/data/addonscripts/concurrent_decorator_class.py"
) )
) )
l = addonmanager.Loader(tctx.master)
sc.load(l)
f1, f2 = tflow.tflow(), tflow.tflow() f1, f2 = tflow.tflow(), tflow.tflow()
tctx.cycle(sc, f1) tctx.cycle(sc, f1)
tctx.cycle(sc, f2) tctx.cycle(sc, f2)