diff options
| -rw-r--r-- | mitmproxy/addonmanager.py | 10 | ||||
| -rw-r--r-- | mitmproxy/addons/script.py | 85 | ||||
| -rw-r--r-- | mitmproxy/test/taddons.py | 11 | ||||
| -rw-r--r-- | setup.py | 1 | ||||
| -rw-r--r-- | test/examples/test_examples.py | 118 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_script.py | 61 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/addon.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/concurrent_decorator_err.py | 2 | ||||
| -rw-r--r-- | 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 b6d7adb6..e763d914 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 bda823b4..5099e62c 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 3dbccba2..471c9c31 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 @@ -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 4b691df2..4c1631ce 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"<html>Test!</html>") - 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"<html>Test!</html>") + 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")) + + original_host = "example.com" - host_header = Headers(host=original_host) - f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) + host_header = Headers(host=original_host) + 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 - f.request.scheme = "https" - f.request.port = 443 + # Rewrite by reverse proxy mode + f.request.scheme = "https" + 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.port == 80 + assert f.request.scheme == "http" + assert f.request.port == 80 - assert f.request.headers["Host"] == original_host + assert f.request.headers["Host"] == original_host diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 4a86fad2..859d99f9 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 42e28a93..8bd25808 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 7bc28182..4f80e98a 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 d24f96a2..ceff9fb9 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) |