From a3a22fba337fc4ac750b8c18663233920a0d646b Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 13:51:00 +1200 Subject: [PATCH 01/13] First-order integration of scripts addon --- mitmproxy/builtins/__init__.py | 2 + mitmproxy/builtins/script.py | 156 ++++++++++++++++++ mitmproxy/controller.py | 12 +- mitmproxy/dump.py | 8 - mitmproxy/exceptions.py | 4 + mitmproxy/flow/master.py | 98 +---------- test/mitmproxy/builtins/test_script.py | 136 +++++++++++++++ .../data/addonscripts/duplicate_flow.py | 6 + test/mitmproxy/data/addonscripts/error.py | 7 + test/mitmproxy/data/addonscripts/recorder.py | 18 ++ .../data/addonscripts/stream_modify.py | 8 + .../data/addonscripts/tcp_stream_modify.py | 5 + test/mitmproxy/test_dump.py | 4 +- test/mitmproxy/test_flow.py | 64 +------ test/mitmproxy/test_script.py | 13 -- test/mitmproxy/test_server.py | 22 ++- 16 files changed, 379 insertions(+), 184 deletions(-) create mode 100644 mitmproxy/builtins/script.py create mode 100644 test/mitmproxy/builtins/test_script.py create mode 100644 test/mitmproxy/data/addonscripts/duplicate_flow.py create mode 100644 test/mitmproxy/data/addonscripts/error.py create mode 100644 test/mitmproxy/data/addonscripts/recorder.py create mode 100644 test/mitmproxy/data/addonscripts/stream_modify.py create mode 100644 test/mitmproxy/data/addonscripts/tcp_stream_modify.py delete mode 100644 test/mitmproxy/test_script.py diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 8021c20ff..6b357902a 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -4,6 +4,7 @@ from mitmproxy.builtins import anticache from mitmproxy.builtins import anticomp from mitmproxy.builtins import stickyauth from mitmproxy.builtins import stickycookie +from mitmproxy.builtins import script from mitmproxy.builtins import stream @@ -13,5 +14,6 @@ def default_addons(): anticomp.AntiComp(), stickyauth.StickyAuth(), stickycookie.StickyCookie(), + script.ScriptLoader(), stream.Stream(), ] diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py new file mode 100644 index 000000000..015adef92 --- /dev/null +++ b/mitmproxy/builtins/script.py @@ -0,0 +1,156 @@ +from __future__ import absolute_import, print_function, division + +import contextlib +import os +import shlex +import sys +import traceback +import copy + +from mitmproxy import exceptions +from mitmproxy import controller +from mitmproxy import ctx + + +import watchdog.events +# The OSX reloader in watchdog 0.8.3 breaks when unobserving paths. +# We use the PollingObserver instead. +if sys.platform == 'darwin': # pragma: no cover + from watchdog.observers.polling import PollingObserver as Observer +else: + from watchdog.observers import Observer + + +def parse_command(command): + """ + Returns a (path, args) tuple. + """ + if not command or not command.strip(): + raise exceptions.AddonError("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.AddonError( + ("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.AddonError("Not a file: %s" % args[0]) + return args[0], args[1:] + + +@contextlib.contextmanager +def scriptenv(path, args): + oldargs = sys.argv + sys.argv = [path] + args + script_dir = os.path.dirname(os.path.abspath(path)) + sys.path.append(script_dir) + try: + yield + except Exception: + _, _, tb = sys.exc_info() + scriptdir = os.path.dirname(os.path.abspath(path)) + for i, s in enumerate(reversed(traceback.extract_tb(tb))): + if not os.path.abspath(s[0]).startswith(scriptdir): + break + else: + tb = tb.tb_next + ctx.log.warn("".join(traceback.format_tb(tb))) + finally: + sys.argv = oldargs + sys.path.pop() + + +def load_script(path, args): + ns = {'__file__': os.path.abspath(path)} + with scriptenv(path, args): + with open(path, "rb") as f: + code = compile(f.read(), path, 'exec') + exec(code, ns, ns) + return ns + + +class ReloadHandler(watchdog.events.FileSystemEventHandler): + def __init__(self, callback, master, options): + self.callback = callback + self.master, self.options = master, options + + def on_modified(self, event): + self.callback(self.master, self.options) + + def on_created(self, event): + self.callback(self.master, self.options) + + +class Script: + """ + An addon that manages a single script. + """ + def __init__(self, command): + self.name = command + + self.command = command + self.path, self.args = parse_command(command) + self.ns = None + self.observer = None + + for i in controller.Events: + def mkprox(): + evt = i + + def prox(*args, **kwargs): + self.run(evt, *args, **kwargs) + return prox + setattr(self, i, mkprox()) + + def run(self, name, *args, **kwargs): + # It's possible for ns to be un-initialised if we failed during + # configure + if self.ns is not None: + func = self.ns.get(name) + if func: + with scriptenv(self.path, self.args): + func(*args, **kwargs) + + def reload(self, master, options): + with master.handlecontext(): + self.ns = None + self.configure(options) + + def configure(self, options): + if not self.observer: + self.observer = Observer() + # Bind the handler to the real underlying master object + self.observer.schedule( + ReloadHandler( + self.reload, + ctx.master, + copy.copy(options), + ), + os.path.dirname(self.path) or "." + ) + self.observer.start() + if not self.ns: + self.ns = load_script(self.path, self.args) + self.run("configure", options) + + +class ScriptLoader(): + """ + An addon that manages loading scripts from options. + """ + def configure(self, options): + for s in options.scripts or []: + if not ctx.master.addons.has_addon(s): + ctx.log.info("Loading script: %s" % s) + sc = Script(s) + ctx.master.addons.add(sc) + for a in ctx.master.addons.chain: + if isinstance(a, Script): + if a.name not in options.scripts or []: + ctx.master.addons.remove(a) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 2f0c8bf29..464842b68 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -44,7 +44,17 @@ class Log(object): def __call__(self, text, level="info"): self.master.add_event(text, level) - # We may want to add .log(), .warn() etc. here at a later point in time + def debug(self, txt): + self(txt, "debug") + + def info(self, txt): + self(txt, "info") + + def warn(self, txt): + self(txt, "warn") + + def error(self, txt): + self(txt, "error") class Master(object): diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 274e01f35..999a709ae 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -93,13 +93,6 @@ class DumpMaster(flow.FlowMaster): not options.keepserving ) - scripts = options.scripts or [] - for command in scripts: - try: - self.load_script(command, use_reloader=True) - except exceptions.ScriptException as e: - raise DumpError(str(e)) - if options.rfile: try: self.load_flows_file(options.rfile) @@ -335,6 +328,5 @@ class DumpMaster(flow.FlowMaster): def run(self): # pragma: no cover if self.options.rfile and not self.options.keepserving: - self.unload_scripts() # make sure to trigger script unload events. return super(DumpMaster, self).run() diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 282784b6d..3b41fe1ca 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -99,3 +99,7 @@ class ControlException(ProxyException): class OptionsError(Exception): pass + + +class AddonError(Exception): + pass diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 27ceee87b..dbb19ed90 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -9,7 +9,6 @@ import netlib.exceptions from mitmproxy import controller from mitmproxy import exceptions from mitmproxy import models -from mitmproxy import script from mitmproxy.flow import io from mitmproxy.flow import modules from mitmproxy.onboarding import app @@ -35,8 +34,6 @@ class FlowMaster(controller.Master): self.server_playback = None # type: Optional[modules.ServerPlaybackState] self.client_playback = None # type: Optional[modules.ClientPlaybackState] self.kill_nonreplay = False - self.scripts = [] # type: List[script.Script] - self.pause_scripts = False self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.refresh_server_playback = False @@ -60,44 +57,6 @@ class FlowMaster(controller.Master): level: debug, info, error """ - def unload_scripts(self): - for s in self.scripts[:]: - self.unload_script(s) - - def unload_script(self, script_obj): - try: - script_obj.unload() - except script.ScriptException as e: - self.add_event("Script error:\n" + str(e), "error") - script.reloader.unwatch(script_obj) - self.scripts.remove(script_obj) - - def load_script(self, command, use_reloader=False): - """ - Loads a script. - - Raises: - ScriptException - """ - s = script.Script(command) - s.load() - if use_reloader: - s.reply = controller.DummyReply() - script.reloader.watch(s, lambda: self.event_queue.put(("script_change", s))) - self.scripts.append(s) - - def _run_single_script_hook(self, script_obj, name, *args, **kwargs): - if script_obj and not self.pause_scripts: - try: - script_obj.run(name, *args, **kwargs) - except script.ScriptException as e: - self.add_event("Script error:\n{}".format(e), "error") - - def run_scripts(self, name, msg): - for script_obj in self.scripts: - if not msg.reply.acked: - self._run_single_script_hook(script_obj, name, msg) - def get_ignore_filter(self): return self.server.config.check_ignore.patterns @@ -298,11 +257,11 @@ class FlowMaster(controller.Master): if not pb and self.kill_nonreplay: f.kill(self) - def replay_request(self, f, block=False, run_scripthooks=True): + def replay_request(self, f, block=False): """ Returns None if successful, or error message if not. """ - if f.live and run_scripthooks: + if f.live: return "Can't replay live request." if f.intercepted: return "Can't replay while intercepting..." @@ -319,7 +278,7 @@ class FlowMaster(controller.Master): rt = http_replay.RequestReplayThread( self.server.config, f, - self.event_queue if run_scripthooks else False, + self.event_queue, self.should_exit ) rt.start() # pragma: no cover @@ -332,28 +291,27 @@ class FlowMaster(controller.Master): @controller.handler def clientconnect(self, root_layer): - self.run_scripts("clientconnect", root_layer) + pass @controller.handler def clientdisconnect(self, root_layer): - self.run_scripts("clientdisconnect", root_layer) + pass @controller.handler def serverconnect(self, server_conn): - self.run_scripts("serverconnect", server_conn) + pass @controller.handler def serverdisconnect(self, server_conn): - self.run_scripts("serverdisconnect", server_conn) + pass @controller.handler def next_layer(self, top_layer): - self.run_scripts("next_layer", top_layer) + pass @controller.handler def error(self, f): self.state.update_flow(f) - self.run_scripts("error", f) if self.client_playback: self.client_playback.clear(f) return f @@ -381,8 +339,6 @@ class FlowMaster(controller.Master): self.setheaders.run(f) if not f.reply.acked: self.process_new_request(f) - if not f.reply.acked: - self.run_scripts("request", f) return f @controller.handler @@ -393,7 +349,6 @@ class FlowMaster(controller.Master): except netlib.exceptions.HttpException: f.reply.kill() return - self.run_scripts("responseheaders", f) return f @controller.handler @@ -404,7 +359,6 @@ class FlowMaster(controller.Master): self.replacehooks.run(f) if not f.reply.acked: self.setheaders.run(f) - self.run_scripts("response", f) if not f.reply.acked: if self.client_playback: self.client_playback.clear(f) @@ -416,46 +370,15 @@ class FlowMaster(controller.Master): def handle_accept_intercept(self, f): self.state.update_flow(f) - @controller.handler - def script_change(self, s): - """ - Handle a script whose contents have been changed on the file system. - - Args: - s (script.Script): the changed script - - Returns: - True, if reloading was successful. - False, otherwise. - """ - ok = True - # We deliberately do not want to fail here. - # In the worst case, we have an "empty" script object. - try: - s.unload() - except script.ScriptException as e: - ok = False - self.add_event('Error reloading "{}":\n{}'.format(s.path, e), 'error') - try: - s.load() - except script.ScriptException as e: - ok = False - self.add_event('Error reloading "{}":\n{}'.format(s.path, e), 'error') - else: - self.add_event('"{}" reloaded.'.format(s.path), 'info') - return ok - @controller.handler def tcp_open(self, flow): # TODO: This would break mitmproxy currently. # self.state.add_flow(flow) self.active_flows.add(flow) - self.run_scripts("tcp_open", flow) @controller.handler def tcp_message(self, flow): - # type: (TCPFlow) -> None - self.run_scripts("tcp_message", flow) + pass @controller.handler def tcp_error(self, flow): @@ -463,13 +386,10 @@ class FlowMaster(controller.Master): repr(flow.server_conn.address), flow.error ), "info") - self.run_scripts("tcp_error", flow) @controller.handler def tcp_close(self, flow): self.active_flows.discard(flow) - self.run_scripts("tcp_close", flow) def shutdown(self): super(FlowMaster, self).shutdown() - self.unload_scripts() diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py new file mode 100644 index 000000000..d33661891 --- /dev/null +++ b/test/mitmproxy/builtins/test_script.py @@ -0,0 +1,136 @@ +import time + +from mitmproxy.builtins import script +from mitmproxy import exceptions +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy import options + +from .. import tutils, mastertest + + +class TestParseCommand: + def test_empty_command(self): + with tutils.raises(exceptions.AddonError): + script.parse_command("") + + with tutils.raises(exceptions.AddonError): + 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("data/scripts/a.py") == ("data/scripts/a.py", []) + assert script.parse_command("data/scripts/a.py foo bar") == ("data/scripts/a.py", ["foo", "bar"]) + assert script.parse_command("data/scripts/a.py 'foo bar'") == ("data/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("data\\scripts\\a.py") == ("data\\scripts\\a.py", []) + assert script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", 'foo \\ bar', []) + + +def test_load_script(): + ns = script.load_script( + tutils.test_data.path( + "data/addonscripts/recorder.py" + ), [] + ) + 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() + m = master.FlowMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/recorder.py" + ) + ) + m.addons.add(sc) + assert sc.ns["call_log"] == [("configure", (options.Options(),), {})] + + sc.ns["call_log"] = [] + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + + recf = sc.ns["call_log"][0] + assert recf[0] == "request" + + def test_reload(self): + s = state.State() + m = RecordingMaster(options.Options(), None, s) + with tutils.tmpdir(): + with open("foo.py", "w"): + pass + sc = script.Script("foo.py") + m.addons.add(sc) + + for _ in range(100): + with open("foo.py", "a") as f: + f.write(".") + time.sleep(0.1) + if m.event_log: + return + raise AssertionError("Change event not detected.") + + def test_exception(self): + s = state.State() + m = RecordingMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path("data/addonscripts/error.py") + ) + m.addons.add(sc) + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + assert m.event_log[0][0] == "warn" + + def test_duplicate_flow(self): + s = state.State() + fm = master.FlowMaster(None, None, s) + fm.addons.add( + script.Script( + tutils.test_data.path("data/addonscripts/duplicate_flow.py") + ) + ) + f = tutils.tflow() + fm.request(f) + assert fm.state.flow_count() == 2 + assert not fm.state.view[0].request.is_replay + assert fm.state.view[1].request.is_replay + + +class TestScriptLoader(mastertest.MasterTest): + def test_simple(self): + s = state.State() + o = options.Options(scripts=[]) + m = master.FlowMaster(o, None, s) + sc = script.ScriptLoader() + m.addons.add(sc) + assert len(m.addons) == 1 + o.update( + scripts = [ + tutils.test_data.path("data/addonscripts/recorder.py") + ] + ) + assert len(m.addons) == 2 + o.update(scripts = []) + assert len(m.addons) == 1 diff --git a/test/mitmproxy/data/addonscripts/duplicate_flow.py b/test/mitmproxy/data/addonscripts/duplicate_flow.py new file mode 100644 index 000000000..b466423cd --- /dev/null +++ b/test/mitmproxy/data/addonscripts/duplicate_flow.py @@ -0,0 +1,6 @@ +from mitmproxy import ctx + + +def request(flow): + f = ctx.master.duplicate_flow(flow) + ctx.master.replay_request(f, block=True) diff --git a/test/mitmproxy/data/addonscripts/error.py b/test/mitmproxy/data/addonscripts/error.py new file mode 100644 index 000000000..8ece9fcea --- /dev/null +++ b/test/mitmproxy/data/addonscripts/error.py @@ -0,0 +1,7 @@ + +def mkerr(): + raise ValueError("Error!") + + +def request(flow): + mkerr() diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py new file mode 100644 index 000000000..728203e38 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -0,0 +1,18 @@ +from mitmproxy import controller +from mitmproxy import ctx + +call_log = [] + +# Keep a log of all possible event calls +evts = list(controller.Events) + ["configure"] +for i in evts: + def mkprox(): + evt = i + + def prox(*args, **kwargs): + lg = (evt, args, kwargs) + if evt != "log": + ctx.log.info(str(lg)) + call_log.append(lg) + return prox + globals()[i] = mkprox() diff --git a/test/mitmproxy/data/addonscripts/stream_modify.py b/test/mitmproxy/data/addonscripts/stream_modify.py new file mode 100644 index 000000000..bc6163421 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/stream_modify.py @@ -0,0 +1,8 @@ + +def modify(chunks): + for chunk in chunks: + yield chunk.replace(b"foo", b"bar") + + +def responseheaders(flow): + flow.response.stream = modify diff --git a/test/mitmproxy/data/addonscripts/tcp_stream_modify.py b/test/mitmproxy/data/addonscripts/tcp_stream_modify.py new file mode 100644 index 000000000..af4ccf7e1 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/tcp_stream_modify.py @@ -0,0 +1,5 @@ + +def tcp_message(flow): + message = flow.messages[-1] + if not message.from_client: + message.content = message.content.replace(b"foo", b"bar") diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index 9686be84d..201386e3e 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -245,12 +245,12 @@ class TestDumpMaster(mastertest.MasterTest): assert "XRESPONSE" in ret assert "XCLIENTDISCONNECT" in ret tutils.raises( - dump.DumpError, + exceptions.AddonError, self.mkmaster, None, scripts=["nonexistent"] ) tutils.raises( - dump.DumpError, + exceptions.AddonError, self.mkmaster, None, scripts=["starterr.py"] ) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 1a07f74d3..c58a97038 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -5,7 +5,7 @@ import netlib.utils from netlib.http import Headers from mitmproxy import filt, controller, flow from mitmproxy.contrib import tnetstring -from mitmproxy.exceptions import FlowReadException, ScriptException +from mitmproxy.exceptions import FlowReadException from mitmproxy.models import Error from mitmproxy.models import Flow from mitmproxy.models import HTTPFlow @@ -674,21 +674,6 @@ class TestSerialize: class TestFlowMaster: - def test_load_script(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - - fm.load_script(tutils.test_data.path("data/scripts/a.py")) - fm.load_script(tutils.test_data.path("data/scripts/a.py")) - fm.unload_scripts() - with tutils.raises(ScriptException): - fm.load_script("nonexistent") - try: - fm.load_script(tutils.test_data.path("data/scripts/starterr.py")) - except ScriptException as e: - assert "ValueError" in str(e) - assert len(fm.scripts) == 0 - def test_getset_ignore(self): p = mock.Mock() p.config.check_ignore = HostMatcher() @@ -708,51 +693,7 @@ class TestFlowMaster: assert "intercepting" in fm.replay_request(f) f.live = True - assert "live" in fm.replay_request(f, run_scripthooks=True) - - def test_script_reqerr(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/reqerr.py")) - f = tutils.tflow() - fm.clientconnect(f.client_conn) - assert fm.request(f) - - def test_script(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - f = tutils.tflow(resp=True) - - f.client_conn.acked = False - fm.clientconnect(f.client_conn) - assert fm.scripts[0].ns["log"][-1] == "clientconnect" - f.server_conn.acked = False - fm.serverconnect(f.server_conn) - assert fm.scripts[0].ns["log"][-1] == "serverconnect" - f.reply.acked = False - fm.request(f) - assert fm.scripts[0].ns["log"][-1] == "request" - f.reply.acked = False - fm.response(f) - assert fm.scripts[0].ns["log"][-1] == "response" - # load second script - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - assert len(fm.scripts) == 2 - f.server_conn.reply.acked = False - fm.clientdisconnect(f.server_conn) - assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" - assert fm.scripts[1].ns["log"][-1] == "clientdisconnect" - - # unload first script - fm.unload_scripts() - assert len(fm.scripts) == 0 - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - - f.error = tutils.terr() - f.reply.acked = False - fm.error(f) - assert fm.scripts[0].ns["log"][-1] == "error" + assert "live" in fm.replay_request(f) def test_duplicate_flow(self): s = flow.State() @@ -789,7 +730,6 @@ class TestFlowMaster: f.error.reply = controller.DummyReply() fm.error(f) - fm.load_script(tutils.test_data.path("data/scripts/a.py")) fm.shutdown() def test_client_playback(self): diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py deleted file mode 100644 index 1e8220f10..000000000 --- a/test/mitmproxy/test_script.py +++ /dev/null @@ -1,13 +0,0 @@ -from mitmproxy import flow -from . import tutils - - -def test_duplicate_flow(): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/duplicate_flow.py")) - f = tutils.tflow() - fm.request(f) - assert fm.state.flow_count() == 2 - assert not fm.state.view[0].request.is_replay - assert fm.state.view[1].request.is_replay diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 9dd8b79c2..b1ca6910c 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -13,6 +13,7 @@ from netlib.http import authentication, http1 from netlib.tutils import raises from pathod import pathoc, pathod +from mitmproxy.builtins import script from mitmproxy import controller from mitmproxy.proxy.config import HostMatcher from mitmproxy.models import Error, HTTPResponse, HTTPFlow @@ -287,10 +288,13 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin): self.master.set_stream_large_bodies(None) def test_stream_modify(self): - self.master.load_script(tutils.test_data.path("data/scripts/stream_modify.py")) + s = script.Script( + tutils.test_data.path("data/addonscripts/stream_modify.py") + ) + self.master.addons.add(s) d = self.pathod('200:b"foo"') - assert d.content == b"bar" - self.master.unload_scripts() + assert d.content == "bar" + self.master.addons.remove(s) class TestHTTPAuth(tservers.HTTPProxyTest): @@ -512,15 +516,15 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): ssl = False def test_tcp_stream_modify(self): - self.master.load_script(tutils.test_data.path("data/scripts/tcp_stream_modify.py")) - + s = script.Script( + tutils.test_data.path("data/addonscripts/tcp_stream_modify.py") + ) + self.master.addons.add(s) self._tcpproxy_on() d = self.pathod('200:b"foo"') self._tcpproxy_off() - - assert d.content == b"bar" - - self.master.unload_scripts() + assert d.content == "bar" + self.master.addons.remove(s) class TestTransparentSSL(tservers.TransparentProxyTest, CommonMixin, TcpMixin): From a6821aad8e9296640c3efd4275e8922dd7c6e43b Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 14:39:07 +1200 Subject: [PATCH 02/13] 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] From b94f5fd361af6255ad4d3c7a88b9a21868736dea Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 16:20:27 +1200 Subject: [PATCH 03/13] Convert examples and example tests for new-style scripts Remove the test that just loads all the example scripts for now - it's a very low-value test, and we need to think of something better. --- examples/custom_contentviews.py | 2 +- examples/filt.py | 2 +- examples/flowwriter.py | 2 +- examples/har_extractor.py | 2 +- examples/iframe_injector.py | 2 +- examples/modify_response_body.py | 2 +- examples/proxapp.py | 2 +- examples/stub.py | 6 +- examples/tls_passthrough.py | 2 +- mitmproxy/controller.py | 1 + netlib/utils.py | 7 + test/mitmproxy/script/test_reloader.py | 34 ----- test/mitmproxy/script/test_script.py | 83 ---------- test/mitmproxy/test_examples.py | 204 +++++++++++-------------- 14 files changed, 105 insertions(+), 246 deletions(-) delete mode 100644 test/mitmproxy/script/test_reloader.py delete mode 100644 test/mitmproxy/script/test_script.py diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py index 5a63e2a05..b10d936f6 100644 --- a/examples/custom_contentviews.py +++ b/examples/custom_contentviews.py @@ -62,7 +62,7 @@ class ViewPigLatin(contentviews.View): pig_view = ViewPigLatin() -def start(): +def configure(options): contentviews.add(pig_view) diff --git a/examples/filt.py b/examples/filt.py index 21744edd8..102d1274d 100644 --- a/examples/filt.py +++ b/examples/filt.py @@ -6,7 +6,7 @@ from mitmproxy import filt state = {} -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") state["filter"] = filt.parse(sys.argv[1]) diff --git a/examples/flowwriter.py b/examples/flowwriter.py index 07c7ca200..d8fbc1f44 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -6,7 +6,7 @@ from mitmproxy.flow import FlowWriter state = {} -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 2a69b9afd..23deb43ac 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -61,7 +61,7 @@ class Context(object): context = Context() -def start(): +def configure(options): """ On start we create a HARLog instance. You will have to adapt this to suit your actual needs of HAR generation. As it will probably be diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py index 70247d311..40934dd39 100644 --- a/examples/iframe_injector.py +++ b/examples/iframe_injector.py @@ -7,7 +7,7 @@ from mitmproxy.models import decoded iframe_url = None -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') global iframe_url diff --git a/examples/modify_response_body.py b/examples/modify_response_body.py index 23ad01511..8b6908a4e 100644 --- a/examples/modify_response_body.py +++ b/examples/modify_response_body.py @@ -8,7 +8,7 @@ from mitmproxy.models import decoded state = {} -def start(): +def configure(options): if len(sys.argv) != 3: raise ValueError('Usage: -s "modify_response_body.py old new"') # You may want to use Python's argparse for more sophisticated argument diff --git a/examples/proxapp.py b/examples/proxapp.py index 2935b5874..095f412a8 100644 --- a/examples/proxapp.py +++ b/examples/proxapp.py @@ -16,7 +16,7 @@ def hello_world(): # Register the app using the magic domain "proxapp" on port 80. Requests to # this domain and port combination will now be routed to the WSGI app instance. -def start(): +def configure(options): mitmproxy.ctx.master.apps.add(app, "proxapp", 80) # SSL works too, but the magic domain needs to be resolvable from the mitmproxy machine due to mitmproxy's design. diff --git a/examples/stub.py b/examples/stub.py index 10b342831..614acee29 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -4,11 +4,11 @@ import mitmproxy """ -def start(): +def configure(options): """ - Called once on script startup, before any other events. + Called once on script startup before any other events, and whenever options changes. """ - mitmproxy.ctx.log("start") + mitmproxy.ctx.log("configure") def clientconnect(root_layer): diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 20e8f9bec..306f55f65 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -113,7 +113,7 @@ class TlsFeedback(TlsLayer): tls_strategy = None -def start(): +def configure(options): global tls_strategy if len(sys.argv) == 2: tls_strategy = ProbabilisticStrategy(float(sys.argv[1])) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 464842b68..bffef58a3 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -32,6 +32,7 @@ Events = frozenset([ "error", "log", + "done", "script_change", ]) diff --git a/netlib/utils.py b/netlib/utils.py index 23c16dc36..9eebf22cb 100644 --- a/netlib/utils.py +++ b/netlib/utils.py @@ -56,6 +56,13 @@ class Data(object): dirname = os.path.dirname(inspect.getsourcefile(m)) self.dirname = os.path.abspath(dirname) + def push(self, subpath): + """ + Change the data object to a path relative to the module. + """ + self.dirname = os.path.join(self.dirname, subpath) + return self + def path(self, path): """ Returns a path to the package data housed at 'path' under this diff --git a/test/mitmproxy/script/test_reloader.py b/test/mitmproxy/script/test_reloader.py deleted file mode 100644 index e33903b95..000000000 --- a/test/mitmproxy/script/test_reloader.py +++ /dev/null @@ -1,34 +0,0 @@ -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.path = "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 deleted file mode 100644 index 48fe65c95..000000000 --- a/test/mitmproxy/script/test_script.py +++ /dev/null @@ -1,83 +0,0 @@ -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("data/scripts/a.py") == ("data/scripts/a.py", []) - assert Script.parse_command("data/scripts/a.py foo bar") == ("data/scripts/a.py", ["foo", "bar"]) - assert Script.parse_command("data/scripts/a.py 'foo bar'") == ("data/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("data\\scripts\\a.py") == ("data\\scripts\\a.py", []) - assert Script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", ['foo \\ bar']) - - -def test_simple(): - with tutils.chdir(tutils.test_data.path("data/scripts")): - s = Script("a.py --var 42") - assert s.path == "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") as s: - s.run("here") - - -def test_script_exception(): - with tutils.chdir(tutils.test_data.path("data/scripts")): - s = Script("syntaxerr.py") - with tutils.raises(ScriptException): - s.load() - - s = Script("starterr.py") - with tutils.raises(ScriptException): - s.load() - - s = Script("a.py") - s.load() - with tutils.raises(ScriptException): - s.load() - - s = Script("a.py") - with tutils.raises(ScriptException): - s.run("here") - - with tutils.raises(ScriptException): - with Script("reqerr.py") as s: - s.run("request", None) - - s = Script("unloaderr.py") - s.load() - with tutils.raises(ScriptException): - s.unload() diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index bdadcd110..9c8edb298 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -1,151 +1,119 @@ -import glob import json -import mock -import os -import sys -from contextlib import contextmanager -from mitmproxy import script +import os.path +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy import options +from mitmproxy import contentviews +from mitmproxy.builtins import script import netlib.utils from netlib import tutils as netutils from netlib.http import Headers -from . import tutils +from . import tutils, mastertest -example_dir = netlib.utils.Data(__name__).path("../../examples") +example_dir = netlib.utils.Data(__name__).push("../../examples") -@contextmanager -def example(command): - command = os.path.join(example_dir, command) - with script.Script(command) as s: - yield s +class ScriptError(Exception): + pass -@mock.patch("mitmproxy.ctx.master") -@mock.patch("mitmproxy.ctx.log") -def test_load_scripts(log, master): - scripts = glob.glob("%s/*.py" % example_dir) - - for f in scripts: - if "har_extractor" in f: - continue - if "flowwriter" in f: - f += " -" - if "iframe_injector" in f: - f += " foo" # one argument required - if "filt" in f: - f += " ~a" - if "modify_response_body" in f: - f += " foo bar" # two arguments required - - s = script.Script(f) - try: - s.load() - except Exception as v: - if "ImportError" not in str(v): - raise - else: - s.unload() +class RaiseMaster(master.FlowMaster): + def add_event(self, e, level): + if level in ("warn", "error"): + raise ScriptError(e) -def test_add_header(): - flow = tutils.tflow(resp=netutils.tresp()) - with example("add_header.py") as ex: - ex.run("response", flow) - assert flow.response.headers["newheader"] == "foo" +def tscript(cmd, args=""): + cmd = example_dir.path(cmd) + " " + args + m = RaiseMaster(options.Options(), None, state.State()) + sc = script.Script(cmd) + m.addons.add(sc) + return m, sc -@mock.patch("mitmproxy.contentviews.remove") -@mock.patch("mitmproxy.contentviews.add") -def test_custom_contentviews(add, remove): - with example("custom_contentviews.py"): - assert add.called - pig = add.call_args[0][0] - _, fmt = pig(b"test!") - assert any(b'esttay!' in val[0][1] for val in fmt) - assert not pig(b"gobbledygook") - assert remove.called +class TestScripts(mastertest.MasterTest): + def test_add_header(self): + m, _ = tscript("add_header.py") + f = tutils.tflow(resp=netutils.tresp()) + self.invoke(m, "response", f) + assert f.response.headers["newheader"] == "foo" + def test_custom_contentviews(self): + m, sc = tscript("custom_contentviews.py") + pig = contentviews.get("pig_latin_HTML") + _, fmt = pig("test!") + assert any('esttay!' in val[0][1] for val in fmt) + assert not pig("gobbledygook") -def test_iframe_injector(): - with tutils.raises(script.ScriptException): - with example("iframe_injector.py"): - pass + def test_iframe_injector(self): + with tutils.raises(ScriptError): + tscript("iframe_injector.py") - flow = tutils.tflow(resp=netutils.tresp(content=b"mitmproxy")) - with example("iframe_injector.py http://example.org/evil_iframe") as ex: - ex.run("response", flow) + m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe") + flow = tutils.tflow(resp=netutils.tresp(content="mitmproxy")) + self.invoke(m, "response", flow) content = flow.response.content - assert b'iframe' in content and b'evil_iframe' in content + assert 'iframe' in content and 'evil_iframe' in content + def test_modify_form(self): + m, sc = tscript("modify_form.py") -def test_modify_form(): - form_header = Headers(content_type="application/x-www-form-urlencoded") - flow = tutils.tflow(req=netutils.treq(headers=form_header)) - with example("modify_form.py") as ex: - ex.run("request", flow) - assert flow.request.urlencoded_form[b"mitmproxy"] == b"rocks" + form_header = Headers(content_type="application/x-www-form-urlencoded") + f = tutils.tflow(req=netutils.treq(headers=form_header)) + self.invoke(m, "request", f) - flow.request.headers["content-type"] = "" - ex.run("request", flow) - assert list(flow.request.urlencoded_form.items()) == [(b"foo", b"bar")] + assert f.request.urlencoded_form["mitmproxy"] == "rocks" + f.request.headers["content-type"] = "" + self.invoke(m, "request", f) + assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] -def test_modify_querystring(): - flow = tutils.tflow(req=netutils.treq(path=b"/search?q=term")) - with example("modify_querystring.py") as ex: - ex.run("request", flow) - assert flow.request.query["mitmproxy"] == "rocks" + def test_modify_querystring(self): + m, sc = tscript("modify_querystring.py") + f = tutils.tflow(req=netutils.treq(path="/search?q=term")) - flow.request.path = "/" - ex.run("request", flow) - assert flow.request.query["mitmproxy"] == "rocks" + self.invoke(m, "request", f) + assert f.request.query["mitmproxy"] == "rocks" + f.request.path = "/" + self.invoke(m, "request", f) + assert f.request.query["mitmproxy"] == "rocks" -def test_modify_response_body(): - with tutils.raises(script.ScriptException): - with example("modify_response_body.py"): - assert True + def test_modify_response_body(self): + with tutils.raises(ScriptError): + tscript("modify_response_body.py") - flow = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy")) - with example("modify_response_body.py mitmproxy rocks") as ex: - assert ex.ns["state"]["old"] == b"mitmproxy" and ex.ns["state"]["new"] == b"rocks" - ex.run("response", flow) - assert flow.response.content == b"I <3 rocks" + m, sc = tscript("modify_response_body.py", "mitmproxy rocks") + f = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy")) + self.invoke(m, "response", f) + assert f.response.content == "I <3 rocks" + def test_redirect_requests(self): + m, sc = tscript("redirect_requests.py") + f = tutils.tflow(req=netutils.treq(host="example.org")) + self.invoke(m, "request", f) + assert f.request.host == "mitmproxy.org" -def test_redirect_requests(): - flow = tutils.tflow(req=netutils.treq(host=b"example.org")) - with example("redirect_requests.py") as ex: - ex.run("request", flow) - assert flow.request.host == "mitmproxy.org" + def test_har_extractor(self): + with tutils.raises(ScriptError): + tscript("har_extractor.py") + with tutils.tmpdir() as tdir: + times = dict( + timestamp_start=746203272, + timestamp_end=746203272, + ) -@mock.patch("mitmproxy.ctx.log") -def test_har_extractor(log): - if sys.version_info >= (3, 0): - with tutils.raises("does not work on Python 3"): - with example("har_extractor.py -"): - pass - return + path = os.path.join(tdir, "file") + m, sc = tscript("har_extractor.py", path) + f = tutils.tflow( + req=netutils.treq(**times), + resp=netutils.tresp(**times) + ) + self.invoke(m, "response", f) + m.addons.remove(sc) - with tutils.raises(script.ScriptException): - with example("har_extractor.py"): - pass - - times = dict( - timestamp_start=746203272, - timestamp_end=746203272, - ) - - flow = tutils.tflow( - req=netutils.treq(**times), - resp=netutils.tresp(**times) - ) - - with example("har_extractor.py -") as ex: - ex.run("response", flow) - - with open(tutils.test_data.path("data/har_extractor.har")) as fp: + fp = open(path, "rb") test_data = json.load(fp) - assert json.loads(ex.ns["context"].HARLog.json()) == test_data["test_response"] + assert len(test_data["log"]["pages"]) == 1 From deffed2196a8d595624998b9fcc8fa4016b41808 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 17:19:33 +1200 Subject: [PATCH 04/13] Script cleanup: editing in console, Python3 compatibility fixes --- mitmproxy/console/grideditor.py | 11 ++++---- mitmproxy/console/master.py | 32 +----------------------- mitmproxy/console/options.py | 16 ++++++------ mitmproxy/console/statusbar.py | 9 +++---- test/mitmproxy/script/test_concurrent.py | 3 +-- test/mitmproxy/test_examples.py | 24 +++++++++++------- test/mitmproxy/test_server.py | 4 +-- 7 files changed, 38 insertions(+), 61 deletions(-) diff --git a/mitmproxy/console/grideditor.py b/mitmproxy/console/grideditor.py index 9fa51ccbe..f304de57a 100644 --- a/mitmproxy/console/grideditor.py +++ b/mitmproxy/console/grideditor.py @@ -6,11 +6,12 @@ import re import urwid +from mitmproxy import exceptions from mitmproxy import filt -from mitmproxy import script -from mitmproxy import utils +from mitmproxy.builtins import script from mitmproxy.console import common from mitmproxy.console import signals +from netlib import strutils from netlib.http import cookies from netlib.http import user_agents @@ -55,7 +56,7 @@ class TextColumn: o = editor.walker.get_current_value() if o is not None: n = editor.master.spawn_editor(o.encode("string-escape")) - n = utils.clean_hanging_newline(n) + n = strutils.clean_hanging_newline(n) editor.walker.set_current_value(n, False) editor.walker._modified() elif key in ["enter"]: @@ -643,8 +644,8 @@ class ScriptEditor(GridEditor): def is_error(self, col, val): try: - script.Script.parse_command(val) - except script.ScriptException as e: + script.parse_command(val) + except exceptions.AddonError as e: return str(e) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index bc373a2b0..64bd9f0af 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -248,23 +248,6 @@ class ConsoleMaster(flow.FlowMaster): if options.server_replay: self.server_playback_path(options.server_replay) - if options.scripts: - for i in options.scripts: - try: - self.load_script(i) - except exceptions.ScriptException as e: - print("Script load error: {}".format(e), file=sys.stderr) - sys.exit(1) - - if options.outfile: - err = self.start_stream_to_path( - options.outfile[0], - options.outfile[1] - ) - if err: - print("Stream file error: {}".format(err), file=sys.stderr) - sys.exit(1) - self.view_stack = [] if options.app: @@ -685,20 +668,7 @@ class ConsoleMaster(flow.FlowMaster): self.refresh_focus() def edit_scripts(self, scripts): - commands = [x[0] for x in scripts] # remove outer array - if commands == [s.command for s in self.scripts]: - return - - self.unload_scripts() - for command in commands: - try: - self.load_script(command) - except exceptions.ScriptException as e: - signals.status_message.send( - message='Error loading "{}".'.format(command) - ) - signals.add_event('Error loading "{}":\n{}'.format(command, e), "error") - signals.update_settings.send(self) + self.options.scripts = [x[0] for x in scripts] def stop_client_playback_prompt(self, a): if a != "n": diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index d363ba747..d8824b056 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -54,7 +54,7 @@ class Options(urwid.WidgetWrap): select.Option( "Scripts", "S", - lambda: master.scripts, + lambda: master.options.scripts, self.scripts ), @@ -160,12 +160,14 @@ class Options(urwid.WidgetWrap): self.master.replacehooks.clear() self.master.set_ignore_filter([]) self.master.set_tcp_filter([]) - self.master.scripts = [] - self.master.options.anticache = False - self.master.options.anticomp = False - self.master.options.stickyauth = None - self.master.options.stickycookie = None + self.master.options.update( + scripts = [], + anticache = False, + anticomp = False, + stickyauth = None, + stickycookie = None + ) self.master.state.default_body_view = contentviews.get("Auto") @@ -234,7 +236,7 @@ class Options(urwid.WidgetWrap): self.master.view_grideditor( grideditor.ScriptEditor( self.master, - [[i.command] for i in self.master.scripts], + [[i] for i in self.master.options.scripts], self.master.edit_scripts ) ) diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index fc41869c3..e7a700a6e 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -218,14 +218,13 @@ class StatusBar(urwid.WidgetWrap): dst.address.host, dst.address.port )) - if self.master.scripts: + if self.master.options.scripts: r.append("[") r.append(("heading_key", "s")) - r.append("cripts:%s]" % len(self.master.scripts)) - # r.append("[lt:%0.3f]"%self.master.looptime) + r.append("cripts:%s]" % len(self.master.options.scripts)) - if self.master.stream: - r.append("[W:%s]" % self.master.stream_path) + if self.master.options.outfile: + r.append("[W:%s]" % self.master.outfile[0]) return r diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index d5243bcbc..92d1153b9 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -1,11 +1,10 @@ -from test.mitmproxy import tutils +from test.mitmproxy import tutils, mastertest 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: diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index 9c8edb298..ef97219cf 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -1,5 +1,6 @@ import json +import sys import os.path from mitmproxy.flow import master from mitmproxy.flow import state @@ -42,19 +43,19 @@ class TestScripts(mastertest.MasterTest): def test_custom_contentviews(self): m, sc = tscript("custom_contentviews.py") pig = contentviews.get("pig_latin_HTML") - _, fmt = pig("test!") - assert any('esttay!' in val[0][1] for val in fmt) - assert not pig("gobbledygook") + _, fmt = pig(b"test!") + assert any(b'esttay!' in val[0][1] for val in fmt) + assert not pig(b"gobbledygook") def test_iframe_injector(self): with tutils.raises(ScriptError): tscript("iframe_injector.py") m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe") - flow = tutils.tflow(resp=netutils.tresp(content="mitmproxy")) + flow = tutils.tflow(resp=netutils.tresp(content=b"mitmproxy")) self.invoke(m, "response", flow) content = flow.response.content - assert 'iframe' in content and 'evil_iframe' in content + assert b'iframe' in content and b'evil_iframe' in content def test_modify_form(self): m, sc = tscript("modify_form.py") @@ -63,11 +64,11 @@ class TestScripts(mastertest.MasterTest): f = tutils.tflow(req=netutils.treq(headers=form_header)) self.invoke(m, "request", f) - assert f.request.urlencoded_form["mitmproxy"] == "rocks" + assert f.request.urlencoded_form[b"mitmproxy"] == b"rocks" f.request.headers["content-type"] = "" self.invoke(m, "request", f) - assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] + assert list(f.request.urlencoded_form.items()) == [(b"foo", b"bar")] def test_modify_querystring(self): m, sc = tscript("modify_querystring.py") @@ -85,9 +86,9 @@ class TestScripts(mastertest.MasterTest): tscript("modify_response_body.py") m, sc = tscript("modify_response_body.py", "mitmproxy rocks") - f = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy")) + f = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy")) self.invoke(m, "response", f) - assert f.response.content == "I <3 rocks" + assert f.response.content == b"I <3 rocks" def test_redirect_requests(self): m, sc = tscript("redirect_requests.py") @@ -96,6 +97,11 @@ class TestScripts(mastertest.MasterTest): assert f.request.host == "mitmproxy.org" def test_har_extractor(self): + if sys.version_info >= (3, 0): + with tutils.raises("does not work on Python 3"): + tscript("har_extractor.py") + return + with tutils.raises(ScriptError): tscript("har_extractor.py") diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index b1ca6910c..a5196daec 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -293,7 +293,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin): ) self.master.addons.add(s) d = self.pathod('200:b"foo"') - assert d.content == "bar" + assert d.content == b"bar" self.master.addons.remove(s) @@ -523,7 +523,7 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): self._tcpproxy_on() d = self.pathod('200:b"foo"') self._tcpproxy_off() - assert d.content == "bar" + assert d.content == b"bar" self.master.addons.remove(s) From 5b2d1c044a0683444f117d8085e29bb613dbbf9d Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 17:22:22 +1200 Subject: [PATCH 05/13] Tighten the tick loop In the past, we consumed from the event queue until we were idle for a certain amount of time (0.1s). This would cause hangs in interactive tools when there was a stream of events, hurting responsiveness. We now wait for a maximum of 0.1s before triggering the tick loop, will be able to reduce this further down the track. --- mitmproxy/controller.py | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index bffef58a3..72f8e001e 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -110,24 +110,21 @@ class Master(object): def tick(self, timeout): changed = False try: - # This endless loop runs until the 'Queue.Empty' - # exception is thrown. - while True: - mtype, obj = self.event_queue.get(timeout=timeout) - if mtype not in Events: - raise exceptions.ControlException("Unknown event %s" % repr(mtype)) - handle_func = getattr(self, mtype) - if not callable(handle_func): - raise exceptions.ControlException("Handler %s not callable" % mtype) - if not handle_func.__dict__.get("__handler"): - raise exceptions.ControlException( - "Handler function %s is not decorated with controller.handler" % ( - handle_func - ) + mtype, obj = self.event_queue.get(timeout=timeout) + if mtype not in Events: + raise exceptions.ControlException("Unknown event %s" % repr(mtype)) + handle_func = getattr(self, mtype) + if not callable(handle_func): + raise exceptions.ControlException("Handler %s not callable" % mtype) + if not handle_func.__dict__.get("__handler"): + raise exceptions.ControlException( + "Handler function %s is not decorated with controller.handler" % ( + handle_func ) - handle_func(obj) - self.event_queue.task_done() - changed = True + ) + handle_func(obj) + self.event_queue.task_done() + changed = True except queue.Empty: pass return changed From a4127fb6d5f026c015525bab7993bf6f33e16f93 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 18:46:07 +1200 Subject: [PATCH 06/13] Add a tick event for addons and scripts, and use it for race-free reload --- mitmproxy/builtins/script.py | 34 +++++++++++++++----------- mitmproxy/controller.py | 3 +++ test/mitmproxy/builtins/test_script.py | 1 + 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index 015adef92..34801ff77 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -4,8 +4,8 @@ import contextlib import os import shlex import sys +import threading import traceback -import copy from mitmproxy import exceptions from mitmproxy import controller @@ -76,15 +76,14 @@ def load_script(path, args): class ReloadHandler(watchdog.events.FileSystemEventHandler): - def __init__(self, callback, master, options): + def __init__(self, callback): self.callback = callback - self.master, self.options = master, options def on_modified(self, event): - self.callback(self.master, self.options) + self.callback() def on_created(self, event): - self.callback(self.master, self.options) + self.callback() class Script: @@ -99,7 +98,10 @@ class Script: self.ns = None self.observer = None - for i in controller.Events: + self.last_options = None + self.should_reload = threading.Event() + + for i in controller.Events - set(["tick"]): def mkprox(): evt = i @@ -117,21 +119,25 @@ class Script: with scriptenv(self.path, self.args): func(*args, **kwargs) - def reload(self, master, options): - with master.handlecontext(): + def reload(self): + self.should_reload.set() + + def tick(self): + if self.should_reload.is_set(): + self.should_reload.clear() self.ns = None - self.configure(options) + ctx.log.info("Reloading script: %s" % self.name) + self.configure(self.last_options) + else: + self.run("tick") def configure(self, options): + self.last_options = options if not self.observer: self.observer = Observer() # Bind the handler to the real underlying master object self.observer.schedule( - ReloadHandler( - self.reload, - ctx.master, - copy.copy(options), - ), + ReloadHandler(self.reload), os.path.dirname(self.path) or "." ) self.observer.start() diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 72f8e001e..d3ae1baa7 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -33,6 +33,7 @@ Events = frozenset([ "error", "log", "done", + "tick", "script_change", ]) @@ -108,6 +109,8 @@ class Master(object): self.shutdown() def tick(self, timeout): + with self.handlecontext(): + self.addons("tick") changed = False try: mtype, obj = self.event_queue.get(timeout=timeout) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 2447c8ea0..394c3f38c 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -78,6 +78,7 @@ class TestScript(mastertest.MasterTest): for _ in range(100): with open("foo.py", "a") as f: f.write(".") + m.addons.invoke_with_context(sc, "tick") time.sleep(0.1) if m.event_log: return From 698af945a49ef34e909f1f5dc1f12552f9e3680b Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 19:58:12 +1200 Subject: [PATCH 07/13] Adjust for new options scheme --- test/mitmproxy/builtins/test_script.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 394c3f38c..7c9787f80 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -4,7 +4,7 @@ from mitmproxy.builtins import script from mitmproxy import exceptions from mitmproxy.flow import master from mitmproxy.flow import state -from mitmproxy import options +from mitmproxy.flow import options from .. import tutils, mastertest From 8fee5db675744a6cc9a3462b7d1261e0b940545c Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Jul 2016 21:33:53 +1200 Subject: [PATCH 08/13] Always use PollingObserver --- mitmproxy/builtins/script.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index 34801ff77..f6cf093f7 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -13,12 +13,7 @@ from mitmproxy import ctx import watchdog.events -# The OSX reloader in watchdog 0.8.3 breaks when unobserving paths. -# We use the PollingObserver instead. -if sys.platform == 'darwin': # pragma: no cover - from watchdog.observers.polling import PollingObserver as Observer -else: - from watchdog.observers import Observer +from watchdog.observers import polling def parse_command(command): @@ -134,7 +129,7 @@ class Script: def configure(self, options): self.last_options = options if not self.observer: - self.observer = Observer() + self.observer = polling.PollingObserver() # Bind the handler to the real underlying master object self.observer.schedule( ReloadHandler(self.reload), From fcc1416ffd82a3497bf17323b4bb467e7e4435f4 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 10:33:51 +1200 Subject: [PATCH 09/13] Fix windows parse_args test --- test/mitmproxy/builtins/test_script.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 7c9787f80..5747912d1 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -35,7 +35,7 @@ class TestParseCommand: def test_parse_windows(self): with tutils.chdir(tutils.test_data.dirname): assert script.parse_command("data\\scripts\\a.py") == ("data\\scripts\\a.py", []) - assert script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", 'foo \\ bar', []) + assert script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", ['foo \\ bar']) def test_load_script(): From 4ba7ce50c6aa7b96325e201f65747a3a6ace1a7a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 13:22:20 +1200 Subject: [PATCH 10/13] Add .start for addons and scripts Also improve error messages, fix various unit tests --- examples/custom_contentviews.py | 2 +- examples/filt.py | 2 +- examples/flowwriter.py | 2 +- examples/har_extractor.py | 2 +- examples/iframe_injector.py | 2 +- examples/modify_response_body.py | 2 +- examples/proxapp.py | 2 +- examples/stub.py | 9 ++++++++ examples/tls_passthrough.py | 2 +- mitmproxy/addons.py | 1 + mitmproxy/builtins/script.py | 29 +++++++++++++++++--------- mitmproxy/controller.py | 3 +++ mitmproxy/script/concurrent.py | 2 +- test/mitmproxy/builtins/test_script.py | 7 +++++-- 14 files changed, 46 insertions(+), 21 deletions(-) diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py index b10d936f6..5a63e2a05 100644 --- a/examples/custom_contentviews.py +++ b/examples/custom_contentviews.py @@ -62,7 +62,7 @@ class ViewPigLatin(contentviews.View): pig_view = ViewPigLatin() -def configure(options): +def start(): contentviews.add(pig_view) diff --git a/examples/filt.py b/examples/filt.py index 102d1274d..21744edd8 100644 --- a/examples/filt.py +++ b/examples/filt.py @@ -6,7 +6,7 @@ from mitmproxy import filt state = {} -def configure(options): +def start(): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") state["filter"] = filt.parse(sys.argv[1]) diff --git a/examples/flowwriter.py b/examples/flowwriter.py index d8fbc1f44..07c7ca200 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -6,7 +6,7 @@ from mitmproxy.flow import FlowWriter state = {} -def configure(options): +def start(): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 23deb43ac..2a69b9afd 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -61,7 +61,7 @@ class Context(object): context = Context() -def configure(options): +def start(): """ On start we create a HARLog instance. You will have to adapt this to suit your actual needs of HAR generation. As it will probably be diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py index 40934dd39..70247d311 100644 --- a/examples/iframe_injector.py +++ b/examples/iframe_injector.py @@ -7,7 +7,7 @@ from mitmproxy.models import decoded iframe_url = None -def configure(options): +def start(): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') global iframe_url diff --git a/examples/modify_response_body.py b/examples/modify_response_body.py index 8b6908a4e..23ad01511 100644 --- a/examples/modify_response_body.py +++ b/examples/modify_response_body.py @@ -8,7 +8,7 @@ from mitmproxy.models import decoded state = {} -def configure(options): +def start(): if len(sys.argv) != 3: raise ValueError('Usage: -s "modify_response_body.py old new"') # You may want to use Python's argparse for more sophisticated argument diff --git a/examples/proxapp.py b/examples/proxapp.py index 095f412a8..2935b5874 100644 --- a/examples/proxapp.py +++ b/examples/proxapp.py @@ -16,7 +16,7 @@ def hello_world(): # Register the app using the magic domain "proxapp" on port 80. Requests to # this domain and port combination will now be routed to the WSGI app instance. -def configure(options): +def start(): mitmproxy.ctx.master.apps.add(app, "proxapp", 80) # SSL works too, but the magic domain needs to be resolvable from the mitmproxy machine due to mitmproxy's design. diff --git a/examples/stub.py b/examples/stub.py index 614acee29..7de4012a9 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -4,6 +4,15 @@ import mitmproxy """ + +def start(): + """ + Called once on script startup before any other events + """ + mitmproxy.ctx.log("start") + + + def configure(options): """ Called once on script startup before any other events, and whenever options changes. diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 306f55f65..20e8f9bec 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -113,7 +113,7 @@ class TlsFeedback(TlsLayer): tls_strategy = None -def configure(options): +def start(): global tls_strategy if len(sys.argv) == 2: tls_strategy = ProbabilisticStrategy(float(sys.argv[1])) diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py index 7ac65a09d..c779aaf8c 100644 --- a/mitmproxy/addons.py +++ b/mitmproxy/addons.py @@ -21,6 +21,7 @@ class Addons(object): def add(self, *addons): self.chain.extend(addons) for i in addons: + self.invoke_with_context(i, "start") self.invoke_with_context(i, "configure", self.master.options) def remove(self, addon): diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index f6cf093f7..bcb756f6e 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -51,22 +51,29 @@ def scriptenv(path, args): _, _, tb = sys.exc_info() scriptdir = os.path.dirname(os.path.abspath(path)) for i, s in enumerate(reversed(traceback.extract_tb(tb))): + tb = tb.tb_next if not os.path.abspath(s[0]).startswith(scriptdir): break - else: - tb = tb.tb_next - ctx.log.warn("".join(traceback.format_tb(tb))) + ctx.log.error("Script error: %s" % "".join(traceback.format_tb(tb))) finally: sys.argv = oldargs sys.path.pop() def load_script(path, args): + with open(path, "rb") as f: + try: + code = compile(f.read(), path, 'exec') + except SyntaxError as e: + ctx.log.error( + "Script error: %s line %s: %s" % ( + e.filename, e.lineno, e.msg + ) + ) + return ns = {'__file__': os.path.abspath(path)} with scriptenv(path, args): - with open(path, "rb") as f: - code = compile(f.read(), path, 'exec') - exec(code, ns, ns) + exec(code, ns, ns) return ns @@ -96,7 +103,7 @@ class Script: self.last_options = None self.should_reload = threading.Event() - for i in controller.Events - set(["tick"]): + for i in controller.Events - set(["start", "configure", "tick"]): def mkprox(): evt = i @@ -120,12 +127,16 @@ class Script: def tick(self): if self.should_reload.is_set(): self.should_reload.clear() - self.ns = None ctx.log.info("Reloading script: %s" % self.name) + self.ns = load_script(self.path, self.args) self.configure(self.last_options) else: self.run("tick") + def start(self): + self.ns = load_script(self.path, self.args) + self.run("start") + def configure(self, options): self.last_options = options if not self.observer: @@ -136,8 +147,6 @@ class Script: os.path.dirname(self.path) or "." ) self.observer.start() - if not self.ns: - self.ns = load_script(self.path, self.args) self.run("configure", options) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index d3ae1baa7..503cdcd3c 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -32,6 +32,9 @@ Events = frozenset([ "error", "log", + + "start", + "configure", "done", "tick", diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index 010a5fa09..0cc0514e7 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -13,7 +13,7 @@ class ScriptThread(basethread.BaseThread): def concurrent(fn): - if fn.__name__ not in controller.Events: + if fn.__name__ not in controller.Events - set(["start", "configure", "tick"]): raise NotImplementedError( "Concurrent decorator not supported for '%s' method." % fn.__name__ ) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 5747912d1..2c2568ed9 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -57,7 +57,10 @@ class TestScript(mastertest.MasterTest): ) ) m.addons.add(sc) - assert sc.ns["call_log"] == [("configure", (options.Options(),), {})] + assert sc.ns["call_log"] == [ + ("start", (), {}), + ("configure", (options.Options(),), {}) + ] sc.ns["call_log"] = [] f = tutils.tflow(resp=True) @@ -93,7 +96,7 @@ class TestScript(mastertest.MasterTest): m.addons.add(sc) f = tutils.tflow(resp=True) self.invoke(m, "request", f) - assert m.event_log[0][0] == "warn" + assert m.event_log[0][0] == "error" def test_duplicate_flow(self): s = state.State() From 92a5076bd01b71e8ed709901d63de9ee37385c3b Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 14:59:48 +1200 Subject: [PATCH 11/13] Adjust concurrent tests for start method --- examples/stub.py | 2 -- test/mitmproxy/script/test_concurrent.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/stub.py b/examples/stub.py index 7de4012a9..e5b4a39ab 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -4,7 +4,6 @@ import mitmproxy """ - def start(): """ Called once on script startup before any other events @@ -12,7 +11,6 @@ def start(): mitmproxy.ctx.log("start") - def configure(options): """ Called once on script startup before any other events, and whenever options changes. diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 92d1153b9..080746e89 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -42,5 +42,5 @@ class TestConcurrent(mastertest.MasterTest): ) ) with m.handlecontext(): - sc.configure(options.Options()) + sc.start() assert "decorator not supported" in m.event_log[0][1] From 917d51bd22ea75408bcf461b09b6cf78c44e1162 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 15:01:35 +1200 Subject: [PATCH 12/13] Fix HAR extractor Thanks to @mhils --- examples/har_extractor.py | 8 +++++--- test/mitmproxy/test_examples.py | 7 ++++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 2a69b9afd..90412ec0f 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -2,7 +2,7 @@ This inline script utilizes harparser.HAR from https://github.com/JustusW/harparser to generate a HAR log object. """ -import mitmproxy +import mitmproxy.ctx import six import sys import pytz @@ -221,9 +221,11 @@ def done(): if context.dump_file == '-': mitmproxy.ctx.log(pprint.pformat(json.loads(json_dump))) elif context.dump_file.endswith('.zhar'): - file(context.dump_file, "w").write(compressed_json_dump) + with open(context.dump_file, "wb") as f: + f.write(compressed_json_dump) else: - file(context.dump_file, "w").write(json_dump) + with open(context.dump_file, "wb") as f: + f.write(json_dump) mitmproxy.ctx.log( "HAR log finished with %s bytes (%s bytes compressed)" % ( len(json_dump), len(compressed_json_dump) diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index ef97219cf..f86463361 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -1,5 +1,6 @@ import json +import six import sys import os.path from mitmproxy.flow import master @@ -112,7 +113,7 @@ class TestScripts(mastertest.MasterTest): ) path = os.path.join(tdir, "file") - m, sc = tscript("har_extractor.py", path) + m, sc = tscript("har_extractor.py", six.moves.shlex_quote(path)) f = tutils.tflow( req=netutils.treq(**times), resp=netutils.tresp(**times) @@ -120,6 +121,6 @@ class TestScripts(mastertest.MasterTest): self.invoke(m, "response", f) m.addons.remove(sc) - fp = open(path, "rb") - test_data = json.load(fp) + with open(path, "rb") as f: + test_data = json.load(f) assert len(test_data["log"]["pages"]) == 1 From c7d0850d8f697915b183f4fafd5ede7df2245569 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 16:35:24 +1200 Subject: [PATCH 13/13] Script cleanups - Preserve script order on config change - Prohibit script duplicates (i.e. identical script + args) - Various cleanups and tweaks --- mitmproxy/builtins/script.py | 47 ++++++++++----- test/mitmproxy/builtins/test_script.py | 62 +++++++++++++++++++- test/mitmproxy/data/addonscripts/recorder.py | 9 ++- 3 files changed, 100 insertions(+), 18 deletions(-) diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index bcb756f6e..14f7dd5f2 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -99,23 +99,25 @@ class Script: self.path, self.args = parse_command(command) self.ns = None self.observer = None + self.dead = False self.last_options = None self.should_reload = threading.Event() - for i in controller.Events - set(["start", "configure", "tick"]): - def mkprox(): - evt = i + for i in controller.Events: + if not hasattr(self, i): + def mkprox(): + evt = i - def prox(*args, **kwargs): - self.run(evt, *args, **kwargs) - return prox - setattr(self, i, mkprox()) + def prox(*args, **kwargs): + self.run(evt, *args, **kwargs) + return prox + setattr(self, i, mkprox()) def run(self, name, *args, **kwargs): # It's possible for ns to be un-initialised if we failed during # configure - if self.ns is not None: + if self.ns is not None and not self.dead: func = self.ns.get(name) if func: with scriptenv(self.path, self.args): @@ -149,18 +151,35 @@ class Script: self.observer.start() self.run("configure", options) + def done(self): + self.run("done") + self.dead = True + class ScriptLoader(): """ An addon that manages loading scripts from options. """ def configure(self, options): - for s in options.scripts or []: - if not ctx.master.addons.has_addon(s): + for s in options.scripts: + if options.scripts.count(s) > 1: + raise exceptions.OptionsError("Duplicate script: %s" % s) + + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script) and a.name not in options.scripts: + ctx.log.info("Un-loading script: %s" % a.name) + ctx.master.addons.remove(a) + + current = {} + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script): + current[a.name] = a + ctx.master.addons.chain.remove(a) + + for s in options.scripts: + if s in current: + ctx.master.addons.chain.append(current[s]) + else: ctx.log.info("Loading script: %s" % s) sc = Script(s) ctx.master.addons.add(sc) - for a in ctx.master.addons.chain: - if isinstance(a, Script): - if a.name not in options.scripts or []: - ctx.master.addons.remove(a) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 2c2568ed9..da60b64cd 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -58,8 +58,8 @@ class TestScript(mastertest.MasterTest): ) m.addons.add(sc) assert sc.ns["call_log"] == [ - ("start", (), {}), - ("configure", (options.Options(),), {}) + ("solo", "start", (), {}), + ("solo", "configure", (options.Options(),), {}) ] sc.ns["call_log"] = [] @@ -67,7 +67,7 @@ class TestScript(mastertest.MasterTest): self.invoke(m, "request", f) recf = sc.ns["call_log"][0] - assert recf[0] == "request" + assert recf[1] == "request" def test_reload(self): s = state.State() @@ -129,3 +129,59 @@ class TestScriptLoader(mastertest.MasterTest): assert len(m.addons) == 2 o.update(scripts = []) assert len(m.addons) == 1 + + def test_dupes(self): + s = state.State() + o = options.Options(scripts=["one", "one"]) + m = master.FlowMaster(o, None, s) + sc = script.ScriptLoader() + tutils.raises(exceptions.OptionsError, m.addons.add, sc) + + def test_order(self): + rec = tutils.test_data.path("data/addonscripts/recorder.py") + + s = state.State() + o = options.Options( + scripts = [ + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + "%s %s" % (rec, "c"), + ] + ) + m = mastertest.RecordingMaster(o, None, s) + sc = script.ScriptLoader() + m.addons.add(sc) + + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'a start'), ('debug', 'a configure'), + ('debug', 'b start'), ('debug', 'b configure'), + ('debug', 'c start'), ('debug', 'c configure') + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "c"), + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c configure'), + ('debug', 'a configure'), + ('debug', 'b configure'), + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "x"), + "%s %s" % (rec, "a"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c done'), + ('debug', 'b done'), + ('debug', 'x start'), + ('debug', 'x configure'), + ('debug', 'a configure'), + ] diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 728203e38..b6ac8d89d 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -1,8 +1,14 @@ from mitmproxy import controller from mitmproxy import ctx +import sys call_log = [] +if len(sys.argv) > 1: + name = sys.argv[1] +else: + name = "solo" + # Keep a log of all possible event calls evts = list(controller.Events) + ["configure"] for i in evts: @@ -10,9 +16,10 @@ for i in evts: evt = i def prox(*args, **kwargs): - lg = (evt, args, kwargs) + lg = (name, evt, args, kwargs) if evt != "log": ctx.log.info(str(lg)) call_log.append(lg) + ctx.log.debug("%s %s" % (name, evt)) return prox globals()[i] = mkprox()