From 7aa208189477f8c5fcd3f7850e1c98fade757f11 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 26 Apr 2017 07:13:36 +1200 Subject: [PATCH] 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. --- mitmproxy/addonmanager.py | 10 +- mitmproxy/addons/script.py | 85 ++++--------- mitmproxy/test/taddons.py | 11 +- setup.py | 1 - test/examples/test_examples.py | 118 ++++++++---------- test/mitmproxy/addons/test_script.py | 61 ++++----- test/mitmproxy/data/addonscripts/addon.py | 2 +- .../addonscripts/concurrent_decorator_err.py | 2 +- test/mitmproxy/script/test_concurrent.py | 12 +- 9 files changed, 118 insertions(+), 184 deletions(-) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index b6d7adb6d..e763d9148 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -129,9 +129,13 @@ class AddonManager: def register(self, addon): """ - Register an addon and all its sub-addons with the manager without - adding it to the chain. This should be used by addons that - dynamically manage addons. Must be called within a current context. + Register an addon, call its load event, and then register all its + sub-addons. This should be used by addons that dynamically manage + 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]): name = _get_name(a) diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index bda823b42..5099e62cc 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -1,15 +1,12 @@ import os import importlib -import threading +import time import sys from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import ctx -import watchdog.events -from watchdog.observers import polling - def load_script(actx, path): if not os.path.exists(path): @@ -28,79 +25,49 @@ def load_script(actx, path): 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: """ An addon that manages a single script. """ + ReloadInterval = 2 + def __init__(self, path): self.name = "scriptmanager:" + path self.path = path self.ns = None - self.observer = None self.last_options = None - self.should_reload = threading.Event() - - def load(self, l): - self.ns = load_script(ctx, self.path) + self.last_load = 0 + self.last_mtime = 0 @property def addons(self): - if self.ns is not None: - return [self.ns] - return [] - - def reload(self): - self.should_reload.set() + return [self.ns] if self.ns else [] def tick(self): - if self.should_reload.is_set(): - self.should_reload.clear() - ctx.log.info("Reloading script: %s" % self.name) - if self.ns: - ctx.master.addons.remove(self.ns) - self.ns = load_script(ctx, self.path) - if self.ns: - # We're already running, so we have to explicitly register and - # configure the addon - ctx.master.addons.register(self.ns) - self.configure(self.last_options, self.last_options.keys()) + if time.time() - self.last_load > self.ReloadInterval: + mtime = os.stat(self.path).st_mtime + if mtime > self.last_mtime: + ctx.log.info("Loading script: %s" % self.name) + if self.ns: + ctx.master.addons.remove(self.ns) + self.ns = load_script(ctx, self.path) + if self.ns: + # We're already running, so we have to explicitly register and + # configure the addon + ctx.master.addons.register(self.ns) + 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): 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: diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 3dbccba2d..471c9c31e 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -112,9 +112,12 @@ class context: ) def script(self, path): + """ + Loads a script from path, and returns the enclosed addon. + """ sc = script.Script(path) loader = addonmanager.Loader(self.master) - sc.load(loader) - for a in addonmanager.traverse(sc.addons): - getattr(a, "load", lambda x: None)(loader) - return sc + self.master.addons.invoke_addon(sc, "load", loader) + self.configure(sc) + self.master.addons.invoke_addon(sc, "tick") + return sc.addons[0] if sc.addons else None diff --git a/setup.py b/setup.py index b6d41b23b..0e9318d0f 100644 --- a/setup.py +++ b/setup.py @@ -80,7 +80,6 @@ setup( "ruamel.yaml>=0.13.2, <0.15", "tornado>=4.3, <4.6", "urwid>=1.3.1, <1.4", - "watchdog>=0.8.3, <0.9", "brotlipy>=0.5.1, <0.7", "sortedcontainers>=1.5.4, <1.6", # transitive from cryptography, we just blacklist here. diff --git a/test/examples/test_examples.py b/test/examples/test_examples.py index 4b691df25..4c1631ce1 100644 --- a/test/examples/test_examples.py +++ b/test/examples/test_examples.py @@ -1,9 +1,4 @@ -from mitmproxy import options 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 tutils from mitmproxy.test import taddons @@ -14,37 +9,20 @@ from ..mitmproxy import tservers 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): def test_add_header(self): - m, _ = tscript("simple/add_header.py") - f = tflow.tflow(resp=tutils.tresp()) - m.addons.handle_lifecycle("response", f) - assert f.response.headers["newheader"] == "foo" + with taddons.context() as tctx: + a = tctx.script(example_dir.path("simple/add_header.py")) + f = tflow.tflow(resp=tutils.tresp()) + a.response(f) + assert f.response.headers["newheader"] == "foo" def test_custom_contentviews(self): - m, sc = tscript("simple/custom_contentview.py") - swapcase = contentviews.get("swapcase") - _, fmt = swapcase(b"Test!") - assert any(b'tEST!' in val[0][1] for val in fmt) + with taddons.context() as tctx: + tctx.script(example_dir.path("simple/custom_contentview.py")) + swapcase = contentviews.get("swapcase") + _, fmt = swapcase(b"Test!") + assert any(b'tEST!' in val[0][1] for val in fmt) def test_iframe_injector(self): with taddons.context() as tctx: @@ -61,57 +39,63 @@ class TestScripts(tservers.MasterTest): assert b'iframe' in content and b'evil_iframe' in content 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") - f = tflow.tflow(req=tutils.treq(headers=form_header)) - m.addons.handle_lifecycle("request", f) + form_header = Headers(content_type="application/x-www-form-urlencoded") + f = tflow.tflow(req=tutils.treq(headers=form_header)) + sc.request(f) - assert f.request.urlencoded_form["mitmproxy"] == "rocks" + assert f.request.urlencoded_form["mitmproxy"] == "rocks" - f.request.headers["content-type"] = "" - m.addons.handle_lifecycle("request", f) - assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] + f.request.headers["content-type"] = "" + sc.request(f) + assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] def test_modify_querystring(self): - m, sc = tscript("simple/modify_querystring.py") - f = tflow.tflow(req=tutils.treq(path="/search?q=term")) + 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")) - m.addons.handle_lifecycle("request", f) - assert f.request.query["mitmproxy"] == "rocks" + sc.request(f) + assert f.request.query["mitmproxy"] == "rocks" - f.request.path = "/" - m.addons.handle_lifecycle("request", f) - assert f.request.query["mitmproxy"] == "rocks" + f.request.path = "/" + sc.request(f) + assert f.request.query["mitmproxy"] == "rocks" def test_redirect_requests(self): - m, sc = tscript("simple/redirect_requests.py") - f = tflow.tflow(req=tutils.treq(host="example.org")) - m.addons.handle_lifecycle("request", f) - assert f.request.host == "mitmproxy.org" + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/redirect_requests.py")) + f = tflow.tflow(req=tutils.treq(host="example.org")) + sc.request(f) + assert f.request.host == "mitmproxy.org" def test_send_reply_from_proxy(self): - m, sc = tscript("simple/send_reply_from_proxy.py") - f = tflow.tflow(req=tutils.treq(host="example.com", port=80)) - m.addons.handle_lifecycle("request", f) - assert f.response.content == b"Hello World" + 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)) + sc.request(f) + assert f.response.content == b"Hello World" def test_dns_spoofing(self): - m, sc = tscript("complex/dns_spoofing.py") - original_host = "example.com" + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("complex/dns_spoofing.py")) - host_header = Headers(host=original_host) - f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) + original_host = "example.com" - m.addons.handle_lifecycle("requestheaders", f) + host_header = Headers(host=original_host) + f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) - # Rewrite by reverse proxy mode - f.request.scheme = "https" - f.request.port = 443 + tctx.master.addons.invoke_addon(sc, "requestheaders", f) - m.addons.handle_lifecycle("request", f) + # Rewrite by reverse proxy mode + f.request.scheme = "https" + f.request.port = 443 - assert f.request.scheme == "http" - assert f.request.port == 80 + tctx.master.addons.invoke_addon(sc, "request", f) - assert f.request.headers["Host"] == original_host + assert f.request.scheme == "http" + assert f.request.port == 80 + + assert f.request.headers["Host"] == original_host diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 4a86fad2c..859d99f98 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -1,7 +1,6 @@ import traceback import sys import time -import watchdog.events import pytest from unittest import mock @@ -16,34 +15,6 @@ from mitmproxy import master 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(): with taddons.context() as tctx: ns = script.load_script( @@ -89,6 +60,8 @@ class TestScript: ) ) tctx.master.addons.add(sc) + tctx.configure(sc) + sc.tick() rec = tctx.master.addons.get("recorder") @@ -107,10 +80,12 @@ class TestScript: f.write("\n") sc = script.Script(str(f)) tctx.configure(sc) - for _ in range(5): - sc.reload() + sc.tick() + for _ in range(3): + sc.last_load, sc.last_mtime = 0, 0 sc.tick() time.sleep(0.1) + tctx.master.has_log("Loading") def test_exception(self): with taddons.context() as tctx: @@ -118,10 +93,12 @@ class TestScript: tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) tctx.master.addons.add(sc) + tctx.configure(sc) + sc.tick() + f = tflow.tflow(resp=True) tctx.master.addons.trigger("request", f) - assert tctx.master.logs[0].level == "error" tctx.master.has_log("ValueError: Error!") tctx.master.has_log("error.py") @@ -133,8 +110,10 @@ class TestScript: ) ) tctx.master.addons.add(sc) + tctx.configure(sc) + sc.tick() assert sc.ns.event_log == [ - 'scriptload', 'addonload' + 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure' ] @@ -207,21 +186,23 @@ class TestScriptLoader: "%s/c.py" % rec, ] ) - + tctx.master.addons.invoke_addon(sc, "tick") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'a load', 'a running', + 'a configure', + 'a tick', 'b load', 'b running', + 'b configure', + 'b tick', 'c load', 'c running', - - 'a configure', - 'b configure', 'c configure', + 'c tick', ] tctx.master.logs = [] @@ -233,6 +214,7 @@ class TestScriptLoader: "%s/b.py" % rec, ] ) + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'c configure', @@ -248,13 +230,16 @@ class TestScriptLoader: "%s/a.py" % rec, ] ) + tctx.master.addons.invoke_addon(sc, "tick") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'c done', 'b done', + 'a configure', 'e load', 'e running', 'e configure', - 'a configure', + 'e tick', + 'a tick', ] diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index 42e28a933..8bd258084 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -14,7 +14,7 @@ class Addon: def configure(options, updated): - event_log.append("addonconfigure") + event_log.append("scriptconfigure") def load(l): diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py index 7bc281820..4f80e98ad 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py @@ -2,5 +2,5 @@ from mitmproxy.script import concurrent @concurrent -def start(opts): +def load(v): pass diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index d24f96a24..ceff9fb92 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -2,10 +2,7 @@ from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons -from mitmproxy import addonmanager from mitmproxy import controller -from mitmproxy.addons import script - import time from .. import tservers @@ -36,25 +33,20 @@ class TestConcurrent(tservers.MasterTest): def test_concurrent_err(self): with taddons.context() as tctx: - sc = script.Script( + tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) assert tctx.master.has_log("decorator not supported") def test_concurrent_class(self): with taddons.context() as tctx: - sc = script.Script( + sc = tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator_class.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) - f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) tctx.cycle(sc, f2)