diff options
author | Thomas Kriechbaumer <Kriechi@users.noreply.github.com> | 2016-03-19 20:33:14 +0100 |
---|---|---|
committer | Thomas Kriechbaumer <Kriechi@users.noreply.github.com> | 2016-03-19 20:33:14 +0100 |
commit | d99194fcccf639008cc140f4cb432b18e1296e57 (patch) | |
tree | 3b26837642025e8112c706649d630e931b6f2cc0 | |
parent | 4b955da94e1dc88e1d4850c6fc3e461be262f789 (diff) | |
parent | 4be9074b492ec49a4ddd163be72ed835e8feb6f3 (diff) | |
download | mitmproxy-d99194fcccf639008cc140f4cb432b18e1296e57.tar.gz mitmproxy-d99194fcccf639008cc140f4cb432b18e1296e57.tar.bz2 mitmproxy-d99194fcccf639008cc140f4cb432b18e1296e57.zip |
Merge pull request #1043 from mitmproxy/better-scripts
Better scripts
-rw-r--r-- | mitmproxy/console/__init__.py | 1 | ||||
-rw-r--r-- | mitmproxy/flow.py | 1 | ||||
-rw-r--r-- | mitmproxy/script/reloader.py | 6 | ||||
-rw-r--r-- | mitmproxy/script/script.py | 12 | ||||
-rw-r--r-- | test/mitmproxy/script/__init__.py | 0 | ||||
-rw-r--r-- | test/mitmproxy/script/test_concurrent.py | 32 | ||||
-rw-r--r-- | test/mitmproxy/script/test_reloader.py | 34 | ||||
-rw-r--r-- | test/mitmproxy/script/test_script.py | 83 | ||||
-rw-r--r-- | test/mitmproxy/scripts/concurrent_decorator.py | 26 | ||||
-rw-r--r-- | test/mitmproxy/test_examples.py | 9 | ||||
-rw-r--r-- | test/mitmproxy/test_script.py | 121 | ||||
-rw-r--r-- | test/mitmproxy/tutils.py | 22 |
12 files changed, 186 insertions, 161 deletions
diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py index 5669be48..2f52d0b8 100644 --- a/mitmproxy/console/__init__.py +++ b/mitmproxy/console/__init__.py @@ -318,6 +318,7 @@ class ConsoleMaster(flow.FlowMaster): try: s = script.Script(command, script.ScriptContext(self)) + s.load() except script.ScriptException as v: signals.status_message.send( message = "Error loading script." diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 8fa84ed8..7de9823d 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -685,6 +685,7 @@ class FlowMaster(controller.Master): """ try: s = script.Script(command, script.ScriptContext(self)) + s.load() except script.ScriptException as v: return v.args[0] if use_reloader: diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py index b4acf51b..99ce7f60 100644 --- a/mitmproxy/script/reloader.py +++ b/mitmproxy/script/reloader.py @@ -1,7 +1,7 @@ import os import sys from watchdog.events import RegexMatchingEventHandler -if sys.platform == 'darwin': +if sys.platform == 'darwin': # pragma: no cover from watchdog.observers.polling import PollingObserver as Observer else: from watchdog.observers import Observer @@ -14,8 +14,8 @@ _observers = {} def watch(script, callback): if script in _observers: raise RuntimeError("Script already observed") - script_dir = os.path.dirname(os.path.abspath(script.args[0])) - script_name = os.path.basename(script.args[0]) + script_dir = os.path.dirname(os.path.abspath(script.filename)) + script_name = os.path.basename(script.filename) event_handler = _ScriptModificationHandler(callback, filename=script_name) observer = Observer() observer.schedule(event_handler, script_dir) diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py index 55778851..edc17d43 100644 --- a/mitmproxy/script/script.py +++ b/mitmproxy/script/script.py @@ -22,7 +22,15 @@ class Script(object): self.args = self.parse_command(command) self.ctx = context self.ns = None + + def __enter__(self): self.load() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_val: + return False # reraise the exception + self.unload() @property def filename(self): @@ -35,7 +43,7 @@ class Script(object): if os.name == "nt": # Windows: escape all backslashes in the path. backslashes = shlex.split(command, posix=False)[0].count("\\") command = command.replace("\\", "\\\\", backslashes) - args = shlex.split(command) + args = shlex.split(command) # pragma: nocover args[0] = os.path.expanduser(args[0]) if not os.path.exists(args[0]): raise ScriptException( @@ -58,7 +66,7 @@ class Script(object): ScriptException on failure """ if self.ns is not None: - self.unload() + raise ScriptException("Script is already loaded") script_dir = os.path.dirname(os.path.abspath(self.args[0])) self.ns = {'__file__': os.path.abspath(self.args[0])} sys.path.append(script_dir) diff --git a/test/mitmproxy/script/__init__.py b/test/mitmproxy/script/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/test/mitmproxy/script/__init__.py diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py new file mode 100644 index 00000000..a9c9e153 --- /dev/null +++ b/test/mitmproxy/script/test_concurrent.py @@ -0,0 +1,32 @@ +from threading import Event + +from mitmproxy.script import Script +from test.mitmproxy import tutils + + +class Dummy: + def __init__(self, reply): + self.reply = reply + + +@tutils.skip_appveyor +def test_concurrent(): + with Script(tutils.test_data.path("scripts/concurrent_decorator.py"), None) as s: + def reply(): + reply.acked.set() + reply.acked = Event() + + f1, f2 = Dummy(reply), Dummy(reply) + s.run("request", f1) + f1.reply() + s.run("request", f2) + f2.reply() + assert f1.reply.acked == reply.acked + assert not reply.acked.is_set() + assert reply.acked.wait(10) + + +def test_concurrent_err(): + s = Script(tutils.test_data.path("scripts/concurrent_decorator_err.py"), None) + with tutils.raises("Concurrent decorator not supported for 'start' method"): + s.load() diff --git a/test/mitmproxy/script/test_reloader.py b/test/mitmproxy/script/test_reloader.py new file mode 100644 index 00000000..0345f6ed --- /dev/null +++ b/test/mitmproxy/script/test_reloader.py @@ -0,0 +1,34 @@ +import mock +from mitmproxy.script.reloader import watch, unwatch +from test.mitmproxy import tutils +from threading import Event + + +def test_simple(): + with tutils.tmpdir(): + with open("foo.py", "w"): + pass + + script = mock.Mock() + script.filename = "foo.py" + + e = Event() + + def _onchange(): + e.set() + + watch(script, _onchange) + with tutils.raises("already observed"): + watch(script, _onchange) + + # Some reloaders don't register a change directly after watching, because they first need to initialize. + # To test if watching works at all, we do repeated writes every 100ms. + for _ in range(100): + with open("foo.py", "a") as f: + f.write(".") + if e.wait(0.1): + break + else: + raise AssertionError("No change detected.") + + unwatch(script) diff --git a/test/mitmproxy/script/test_script.py b/test/mitmproxy/script/test_script.py new file mode 100644 index 00000000..a9b55977 --- /dev/null +++ b/test/mitmproxy/script/test_script.py @@ -0,0 +1,83 @@ +from mitmproxy.script import Script +from mitmproxy.exceptions import ScriptException +from test.mitmproxy import tutils + + +class TestParseCommand: + def test_empty_command(self): + with tutils.raises(ScriptException): + Script.parse_command("") + + with tutils.raises(ScriptException): + Script.parse_command(" ") + + def test_no_script_file(self): + with tutils.raises("not found"): + Script.parse_command("notfound") + + with tutils.tmpdir() as dir: + with tutils.raises("not a file"): + Script.parse_command(dir) + + def test_parse_args(self): + with tutils.chdir(tutils.test_data.dirname): + assert Script.parse_command("scripts/a.py") == ["scripts/a.py"] + assert Script.parse_command("scripts/a.py foo bar") == ["scripts/a.py", "foo", "bar"] + assert Script.parse_command("scripts/a.py 'foo bar'") == ["scripts/a.py", "foo bar"] + + @tutils.skip_not_windows + def test_parse_windows(self): + with tutils.chdir(tutils.test_data.dirname): + assert Script.parse_command("scripts\\a.py") == ["scripts\\a.py"] + assert Script.parse_command("scripts\\a.py 'foo \\ bar'") == ["scripts\\a.py", 'foo \\ bar'] + + +def test_simple(): + with tutils.chdir(tutils.test_data.path("scripts")): + s = Script("a.py --var 42", None) + assert s.filename == "a.py" + assert s.ns is None + + s.load() + assert s.ns["var"] == 42 + + s.run("here") + assert s.ns["var"] == 43 + + s.unload() + assert s.ns is None + + with tutils.raises(ScriptException): + s.run("here") + + with Script("a.py --var 42", None) as s: + s.run("here") + + +def test_script_exception(): + with tutils.chdir(tutils.test_data.path("scripts")): + s = Script("syntaxerr.py", None) + with tutils.raises(ScriptException): + s.load() + + s = Script("starterr.py", None) + with tutils.raises(ScriptException): + s.load() + + s = Script("a.py", None) + s.load() + with tutils.raises(ScriptException): + s.load() + + s = Script("a.py", None) + with tutils.raises(ScriptException): + s.run("here") + + with tutils.raises(ScriptException): + with Script("reqerr.py", None) as s: + s.run("request", None) + + s = Script("unloaderr.py", None) + s.load() + with tutils.raises(ScriptException): + s.unload() diff --git a/test/mitmproxy/scripts/concurrent_decorator.py b/test/mitmproxy/scripts/concurrent_decorator.py index 6651c811..cf3ab512 100644 --- a/test/mitmproxy/scripts/concurrent_decorator.py +++ b/test/mitmproxy/scripts/concurrent_decorator.py @@ -1,32 +1,6 @@ import time from mitmproxy.script import concurrent - -@concurrent -def clientconnect(context, cc): - context.log("clientconnect") - - -@concurrent -def serverconnect(context, sc): - context.log("serverconnect") - - @concurrent def request(context, flow): time.sleep(0.1) - - -@concurrent -def response(context, flow): - context.log("response") - - -@concurrent -def error(context, err): - context.log("error") - - -@concurrent -def clientdisconnect(context, dc): - context.log("clientdisconnect") diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index 803776ac..b560d9a1 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -31,9 +31,8 @@ class DummyContext(object): def example(command): command = os.path.join(example_dir, command) ctx = DummyContext() - s = script.Script(command, ctx) - yield s - s.unload() + with script.Script(command, ctx) as s: + yield s def test_load_scripts(): @@ -52,8 +51,10 @@ def test_load_scripts(): f += " ~a" if "modify_response_body" in f: f += " foo bar" # two arguments required + + s = script.Script(f, script.ScriptContext(tmaster)) try: - s = script.Script(f, script.ScriptContext(tmaster)) # Loads the script file. + s.load() except Exception as v: if "ImportError" not in str(v): raise diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py index b827c623..f321d15c 100644 --- a/test/mitmproxy/test_script.py +++ b/test/mitmproxy/test_script.py @@ -1,27 +1,7 @@ -import os -import time -import mock -from mitmproxy import script, flow +from mitmproxy import flow from . import tutils -def test_simple(): - s = flow.State() - fm = flow.FlowMaster(None, s) - sp = tutils.test_data.path("scripts/a.py") - p = script.Script("%s --var 40" % sp, script.ScriptContext(fm)) - - assert "here" in p.ns - assert p.run("here") == 41 - assert p.run("here") == 42 - - tutils.raises(script.ScriptException, p.run, "errargs") - - # Check reload - p.load() - assert p.run("here") == 41 - - def test_duplicate_flow(): s = flow.State() fm = flow.FlowMaster(None, s) @@ -31,102 +11,3 @@ def test_duplicate_flow(): assert fm.state.flow_count() == 2 assert not fm.state.view[0].request.is_replay assert fm.state.view[1].request.is_replay - - -def test_err(): - s = flow.State() - fm = flow.FlowMaster(None, s) - sc = script.ScriptContext(fm) - - tutils.raises( - "not found", - script.Script, "nonexistent", sc - ) - - tutils.raises( - "not a file", - script.Script, tutils.test_data.path("scripts"), sc - ) - - tutils.raises( - script.ScriptException, - script.Script, tutils.test_data.path("scripts/syntaxerr.py"), sc - ) - - tutils.raises( - script.ScriptException, - script.Script, tutils.test_data.path("scripts/loaderr.py"), sc - ) - - scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), sc) - tutils.raises(script.ScriptException, scr.unload) - - -@tutils.skip_appveyor -def test_concurrent(): - s = flow.State() - fm = flow.FlowMaster(None, s) - fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py")) - - with mock.patch("mitmproxy.controller.DummyReply.__call__") as m: - f1, f2 = tutils.tflow(), tutils.tflow() - t_start = time.time() - fm.handle_request(f1) - f1.reply() - fm.handle_request(f2) - f2.reply() - - # Two instantiations - assert m.call_count == 0 # No calls yet. - assert (time.time() - t_start) < 0.1 - - -def test_concurrent2(): - s = flow.State() - fm = flow.FlowMaster(None, s) - s = script.Script( - tutils.test_data.path("scripts/concurrent_decorator.py"), - script.ScriptContext(fm)) - s.load() - m = mock.Mock() - - class Dummy: - - def __init__(self): - self.response = self - self.error = self - self.reply = m - - t_start = time.time() - - for hook in ("clientconnect", - "serverconnect", - "response", - "error", - "clientconnect"): - d = Dummy() - s.run(hook, d) - d.reply() - while (time.time() - t_start) < 20 and m.call_count <= 5: - if m.call_count == 5: - return - time.sleep(0.001) - assert False - - -def test_concurrent_err(): - s = flow.State() - fm = flow.FlowMaster(None, s) - tutils.raises( - "Concurrent decorator not supported for 'start' method", - script.Script, - tutils.test_data.path("scripts/concurrent_decorator_err.py"), - fm) - - -def test_command_parsing(): - s = flow.State() - fm = flow.FlowMaster(None, s) - absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py")) - s = script.Script(absfilepath, script.ScriptContext(fm)) - assert os.path.isfile(s.args[0]) diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py index 0d65df71..791db6d9 100644 --- a/test/mitmproxy/tutils.py +++ b/test/mitmproxy/tutils.py @@ -26,6 +26,13 @@ def skip_windows(fn): return fn +def skip_not_windows(fn): + if os.name == "nt": + return fn + else: + return _skip_windows + + def _skip_appveyor(*args): raise SkipTest("Skipped on AppVeyor.") @@ -120,14 +127,17 @@ def get_body_line(last_displayed_body, line_nb): @contextmanager +def chdir(dir): + orig_dir = os.getcwd() + os.chdir(dir) + yield + os.chdir(orig_dir) + +@contextmanager def tmpdir(*args, **kwargs): - orig_workdir = os.getcwd() temp_workdir = tempfile.mkdtemp(*args, **kwargs) - os.chdir(temp_workdir) - - yield temp_workdir - - os.chdir(orig_workdir) + with chdir(temp_workdir): + yield temp_workdir shutil.rmtree(temp_workdir) |