From a6821aad8e9296640c3efd4275e8922dd7c6e43b Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 14:39:07 +1200 Subject: [PATCH] Zap old scripts infrastructure, fix concurrency tests --- mitmproxy/flow/master.py | 3 - mitmproxy/script/__init__.py | 6 - mitmproxy/script/reloader.py | 47 ------ mitmproxy/script/script.py | 136 ------------------ test/mitmproxy/builtins/test_script.py | 13 +- .../concurrent_decorator.py | 1 - .../concurrent_decorator_err.py | 0 test/mitmproxy/mastertest.py | 10 ++ test/mitmproxy/script/test_concurrent.py | 43 ++++-- 9 files changed, 43 insertions(+), 216 deletions(-) delete mode 100644 mitmproxy/script/reloader.py delete mode 100644 mitmproxy/script/script.py rename test/mitmproxy/data/{scripts => addonscripts}/concurrent_decorator.py (99%) rename test/mitmproxy/data/{scripts => addonscripts}/concurrent_decorator_err.py (100%) diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index dbb19ed90..aa09e1097 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -390,6 +390,3 @@ class FlowMaster(controller.Master): @controller.handler def tcp_close(self, flow): self.active_flows.discard(flow) - - def shutdown(self): - super(FlowMaster, self).shutdown() diff --git a/mitmproxy/script/__init__.py b/mitmproxy/script/__init__.py index 9a3985ab8..e75f282ae 100644 --- a/mitmproxy/script/__init__.py +++ b/mitmproxy/script/__init__.py @@ -1,11 +1,5 @@ -from . import reloader from .concurrent import concurrent -from .script import Script -from ..exceptions import ScriptException __all__ = [ - "Script", "concurrent", - "ScriptException", - "reloader" ] diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py deleted file mode 100644 index 857d76cd5..000000000 --- a/mitmproxy/script/reloader.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import absolute_import, print_function, division - -import os - -from watchdog.events import RegexMatchingEventHandler - -from watchdog.observers.polling import PollingObserver as Observer -# We occasionally have watchdog errors on Windows, Linux and Mac when using the native observers. -# After reading through the watchdog source code and issue tracker, -# we may want to replace this with a very simple implementation of our own. - -_observers = {} - - -def watch(script, callback): - if script in _observers: - raise RuntimeError("Script already observed") - script_dir = os.path.dirname(os.path.abspath(script.path)) - script_name = os.path.basename(script.path) - event_handler = _ScriptModificationHandler(callback, filename=script_name) - observer = Observer() - observer.schedule(event_handler, script_dir) - observer.start() - _observers[script] = observer - - -def unwatch(script): - observer = _observers.pop(script, None) - if observer: - observer.stop() - observer.join() - - -class _ScriptModificationHandler(RegexMatchingEventHandler): - - def __init__(self, callback, filename='.*'): - - super(_ScriptModificationHandler, self).__init__( - ignore_directories=True, - regexes=['.*' + filename] - ) - self.callback = callback - - def on_modified(self, event): - self.callback() - -__all__ = ["watch", "unwatch"] diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py deleted file mode 100644 index db4909cab..000000000 --- a/mitmproxy/script/script.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -The script object representing mitmproxy inline scripts. -Script objects know nothing about mitmproxy or mitmproxy's API - this knowledge is provided -by the mitmproxy-specific ScriptContext. -""" -# Do not import __future__ here, this would apply transitively to the inline scripts. -from __future__ import absolute_import, print_function, division - -import os -import shlex -import sys -import contextlib - -import six -from typing import List # noqa - -from mitmproxy import exceptions - - -@contextlib.contextmanager -def scriptenv(path, args): - # type: (str, List[str]) -> None - oldargs = sys.argv - script_dir = os.path.dirname(os.path.abspath(path)) - - sys.argv = [path] + args - sys.path.append(script_dir) - try: - yield - finally: - sys.argv = oldargs - sys.path.pop() - - -class Script(object): - """ - Script object representing an inline script. - """ - - def __init__(self, command): - self.command = command - self.path, self.args = self.parse_command(command) - self.ns = None - - def __enter__(self): - self.load() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if exc_val: - return False # re-raise the exception - self.unload() - - @staticmethod - def parse_command(command): - # type: (str) -> Tuple[str,List[str]] - """ - Returns a (path, args) tuple. - """ - if not command or not command.strip(): - raise exceptions.ScriptException("Empty script command.") - # Windows: escape all backslashes in the path. - if os.name == "nt": # pragma: no cover - backslashes = shlex.split(command, posix=False)[0].count("\\") - command = command.replace("\\", "\\\\", backslashes) - args = shlex.split(command) # pragma: no cover - args[0] = os.path.expanduser(args[0]) - if not os.path.exists(args[0]): - raise exceptions.ScriptException( - ("Script file not found: %s.\r\n" - "If your script path contains spaces, " - "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") % - args[0]) - elif os.path.isdir(args[0]): - raise exceptions.ScriptException("Not a file: %s" % args[0]) - return args[0], args[1:] - - def load(self): - """ - Loads an inline script. - - Returns: - The return value of self.run("start", ...) - - Raises: - ScriptException on failure - """ - if self.ns is not None: - raise exceptions.ScriptException("Script is already loaded") - self.ns = {'__file__': os.path.abspath(self.path)} - - with scriptenv(self.path, self.args): - try: - with open(self.path) as f: - code = compile(f.read(), self.path, 'exec') - exec(code, self.ns, self.ns) - except Exception: - six.reraise( - exceptions.ScriptException, - exceptions.ScriptException.from_exception_context(), - sys.exc_info()[2] - ) - return self.run("start") - - def unload(self): - try: - return self.run("done") - finally: - self.ns = None - - def run(self, name, *args, **kwargs): - """ - Runs an inline script hook. - - Returns: - The return value of the method. - None, if the script does not provide the method. - - Raises: - ScriptException if there was an exception. - """ - if self.ns is None: - raise exceptions.ScriptException("Script not loaded.") - f = self.ns.get(name) - if f: - try: - with scriptenv(self.path, self.args): - return f(*args, **kwargs) - except Exception: - six.reraise( - exceptions.ScriptException, - exceptions.ScriptException.from_exception_context(), - sys.exc_info()[2] - ) - else: - return None diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index d33661891..2447c8ea0 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -47,15 +47,6 @@ def test_load_script(): assert ns["configure"] -class RecordingMaster(master.FlowMaster): - def __init__(self, *args, **kwargs): - master.FlowMaster.__init__(self, *args, **kwargs) - self.event_log = [] - - def add_event(self, e, level): - self.event_log.append((level, e)) - - class TestScript(mastertest.MasterTest): def test_simple(self): s = state.State() @@ -77,7 +68,7 @@ class TestScript(mastertest.MasterTest): def test_reload(self): s = state.State() - m = RecordingMaster(options.Options(), None, s) + m = mastertest.RecordingMaster(options.Options(), None, s) with tutils.tmpdir(): with open("foo.py", "w"): pass @@ -94,7 +85,7 @@ class TestScript(mastertest.MasterTest): def test_exception(self): s = state.State() - m = RecordingMaster(options.Options(), None, s) + m = mastertest.RecordingMaster(options.Options(), None, s) sc = script.Script( tutils.test_data.path("data/addonscripts/error.py") ) diff --git a/test/mitmproxy/data/scripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py similarity index 99% rename from test/mitmproxy/data/scripts/concurrent_decorator.py rename to test/mitmproxy/data/addonscripts/concurrent_decorator.py index 162c00f4b..a56c2af1c 100644 --- a/test/mitmproxy/data/scripts/concurrent_decorator.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator.py @@ -1,7 +1,6 @@ import time from mitmproxy.script import concurrent - @concurrent def request(flow): time.sleep(0.1) diff --git a/test/mitmproxy/data/scripts/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py similarity index 100% rename from test/mitmproxy/data/scripts/concurrent_decorator_err.py rename to test/mitmproxy/data/addonscripts/concurrent_decorator_err.py diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 9754d3a99..240f6a730 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -3,6 +3,7 @@ import mock from . import tutils import netlib.tutils +from mitmproxy.flow import master from mitmproxy import flow, proxy, models, controller @@ -39,3 +40,12 @@ class MasterTest: t = tutils.tflow(resp=True) fw.add(t) f.close() + + +class RecordingMaster(master.FlowMaster): + def __init__(self, *args, **kwargs): + master.FlowMaster.__init__(self, *args, **kwargs) + self.event_log = [] + + def add_event(self, e, level): + self.event_log.append((level, e)) diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 57eeca191..d5243bcbc 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -1,28 +1,47 @@ -from mitmproxy.script import Script from test.mitmproxy import tutils from mitmproxy import controller +from mitmproxy.builtins import script +from mitmproxy import options +from mitmproxy.flow import master +from mitmproxy.flow import state import time +from .. import mastertest, tutils class Thing: def __init__(self): self.reply = controller.DummyReply() + self.live = True -@tutils.skip_appveyor -def test_concurrent(): - with Script(tutils.test_data.path("data/scripts/concurrent_decorator.py")) as s: - f1, f2 = Thing(), Thing() - s.run("request", f1) - s.run("request", f2) +class TestConcurrent(mastertest.MasterTest): + @tutils.skip_appveyor + def test_concurrent(self): + s = state.State() + m = master.FlowMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/concurrent_decorator.py" + ) + ) + m.addons.add(sc) + f1, f2 = tutils.tflow(), tutils.tflow() + self.invoke(m, "request", f1) + self.invoke(m, "request", f2) start = time.time() while time.time() - start < 5: if f1.reply.acked and f2.reply.acked: return raise ValueError("Script never acked") - -def test_concurrent_err(): - s = Script(tutils.test_data.path("data/scripts/concurrent_decorator_err.py")) - with tutils.raises("Concurrent decorator not supported for 'start' method"): - s.load() + def test_concurrent_err(self): + s = state.State() + m = mastertest.RecordingMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/concurrent_decorator_err.py" + ) + ) + with m.handlecontext(): + sc.configure(options.Options()) + assert "decorator not supported" in m.event_log[0][1]