diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py index 5b980572a..e75aed863 100644 --- a/mitmproxy/console/__init__.py +++ b/mitmproxy/console/__init__.py @@ -19,7 +19,7 @@ from netlib import tcp from .. import flow, script, contentviews from . import flowlist, flowview, help, window, signals, options from . import grideditor, palettes, statusbar, palettepicker -from ..exceptions import FlowReadException +from ..exceptions import FlowReadException, ScriptException EVENTLOG_SIZE = 500 @@ -229,9 +229,10 @@ class ConsoleMaster(flow.FlowMaster): if options.scripts: for i in options.scripts: - err = self.load_script(i) - if err: - print("Script load error: {}".format(err), file=sys.stderr) + try: + self.load_script(i) + except ScriptException as e: + print("Script load error: {}".format(e), file=sys.stderr) sys.exit(1) if options.outfile: @@ -320,11 +321,11 @@ class ConsoleMaster(flow.FlowMaster): try: s = script.Script(command, script.ScriptContext(self)) s.load() - except script.ScriptException as v: + except script.ScriptException as e: signals.status_message.send( - message = "Error loading script." + message='Error loading "{}".'.format(command) ) - signals.add_event("Error loading script:\n%s" % v.args[0], "error") + signals.add_event('Error loading "{}":\n{}'.format(command, e), "error") return if f.request: @@ -336,13 +337,6 @@ class ConsoleMaster(flow.FlowMaster): s.unload() signals.flow_change.send(self, flow = f) - def set_script(self, command): - if not command: - return - ret = self.load_script(command) - if ret: - signals.status_message.send(message=ret) - def toggle_eventlog(self): self.eventlog = not self.eventlog signals.pop_view_state.send(self) @@ -670,7 +664,13 @@ class ConsoleMaster(flow.FlowMaster): self.unload_scripts() for command in commands: - self.load_script(command) + try: + self.load_script(command) + except 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) def stop_client_playback_prompt(self, a): diff --git a/mitmproxy/console/grideditor.py b/mitmproxy/console/grideditor.py index 597a7e7a9..46ff348e9 100644 --- a/mitmproxy/console/grideditor.py +++ b/mitmproxy/console/grideditor.py @@ -642,8 +642,8 @@ class ScriptEditor(GridEditor): def is_error(self, col, val): try: script.Script.parse_command(val) - except script.ScriptException as v: - return str(v) + except script.ScriptException as e: + return str(e) class HostPatternEditor(GridEditor): diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index d224f1aa0..0e0ccc62e 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -7,7 +7,7 @@ import itertools from netlib import tcp import netlib.utils from . import flow, filt, contentviews -from .exceptions import ContentViewException, FlowReadException +from .exceptions import ContentViewException, FlowReadException, ScriptException class DumpError(Exception): @@ -125,9 +125,10 @@ class DumpMaster(flow.FlowMaster): scripts = options.scripts or [] for command in scripts: - err = self.load_script(command, use_reloader=True) - if err: - raise DumpError(err) + try: + self.load_script(command, use_reloader=True) + except ScriptException as e: + raise DumpError(str(e)) if options.rfile: try: diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 86bf75aec..52fd36d11 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -7,6 +7,10 @@ See also: http://lucumr.pocoo.org/2014/10/16/on-error-handling/ """ from __future__ import (absolute_import, print_function, division) +import traceback + +import sys + class ProxyException(Exception): """ @@ -59,7 +63,23 @@ class ReplayException(ProxyException): class ScriptException(ProxyException): - pass + @classmethod + def from_exception_context(cls, cut_tb=1): + """ + Must be called while the current stack handles an exception. + + Args: + cut_tb: remove N frames from the stack trace to hide internal calls. + """ + exc_type, exc_value, exc_traceback = sys.exc_info() + + while cut_tb > 0: + exc_traceback = exc_traceback.tb_next + cut_tb -= 1 + + tb = "".join(traceback.format_exception(exc_type, exc_value, exc_traceback)) + + return cls(tb) class FlowReadException(ProxyException): diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 9292e76ab..c63c1efa6 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -695,14 +695,13 @@ class FlowMaster(controller.ServerMaster): def load_script(self, command, use_reloader=False): """ - Loads a script. Returns an error description if something went - wrong. + Loads a script. + + Raises: + ScriptException """ - try: - s = script.Script(command, script.ScriptContext(self)) - s.load() - except script.ScriptException as e: - return traceback.format_exc(e) + s = script.Script(command, script.ScriptContext(self)) + s.load() if use_reloader: script.reloader.watch(s, lambda: self.event_queue.put(("script_change", s))) self.scripts.append(s) @@ -712,7 +711,7 @@ class FlowMaster(controller.ServerMaster): try: script_obj.run(name, *args, **kwargs) except script.ScriptException as e: - self.add_event("Script error:\n" + str(e), "error") + self.add_event("Script error:\n{}".format(e), "error") def run_script_hook(self, name, *args, **kwargs): for script_obj in self.scripts: @@ -1069,12 +1068,12 @@ class FlowMaster(controller.ServerMaster): s.unload() except script.ScriptException as e: ok = False - self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)), 'error') + self.add_event('Error reloading "{}":\n{}'.format(s.filename, e), 'error') try: s.load() except script.ScriptException as e: ok = False - self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)), 'error') + self.add_event('Error reloading "{}":\n{}'.format(s.filename, e), 'error') else: self.add_event('"{}" reloaded.'.format(s.filename), 'info') return ok diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py index 4d0f73fa9..484025b4a 100644 --- a/mitmproxy/script/script.py +++ b/mitmproxy/script/script.py @@ -79,10 +79,10 @@ class Script(object): with open(self.filename) as f: code = compile(f.read(), self.filename, 'exec') exec (code, self.ns, self.ns) - except Exception as e: + except Exception: six.reraise( ScriptException, - ScriptException(str(e)), + ScriptException.from_exception_context(), sys.exc_info()[2] ) finally: @@ -113,10 +113,10 @@ class Script(object): if f: try: return f(self.ctx, *args, **kwargs) - except Exception as e: + except Exception: six.reraise( ScriptException, - ScriptException(str(e)), + ScriptException.from_exception_context(), sys.exc_info()[2] ) else: diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index b560d9a17..c401a6b93 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -106,8 +106,8 @@ def test_modify_querystring(): def test_modify_response_body(): with tutils.raises(script.ScriptException): - with example("modify_response_body.py") as ex: - pass + with example("modify_response_body.py"): + assert True flow = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy")) with example("modify_response_body.py mitmproxy rocks") as ex: @@ -125,7 +125,7 @@ def test_redirect_requests(): def test_har_extractor(): with tutils.raises(script.ScriptException): - with example("har_extractor.py") as ex: + with example("har_extractor.py"): pass times = dict( diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 145e91cf2..84a25689c 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -7,7 +7,7 @@ import netlib.utils from netlib import odict from netlib.http import Headers from mitmproxy import filt, controller, tnetstring, flow -from mitmproxy.exceptions import FlowReadException +from mitmproxy.exceptions import FlowReadException, ScriptException from mitmproxy.models import Error from mitmproxy.models import Flow from mitmproxy.models import HTTPFlow @@ -747,12 +747,16 @@ class TestFlowMaster: def test_load_script(self): s = flow.State() fm = flow.FlowMaster(None, s) - assert not fm.load_script(tutils.test_data.path("scripts/a.py")) - assert not fm.load_script(tutils.test_data.path("scripts/a.py")) - assert not fm.unload_scripts() - assert fm.load_script("nonexistent") - assert "ValueError" in fm.load_script( - tutils.test_data.path("scripts/starterr.py")) + + fm.load_script(tutils.test_data.path("scripts/a.py")) + fm.load_script(tutils.test_data.path("scripts/a.py")) + fm.unload_scripts() + with tutils.raises(ScriptException): + fm.load_script("nonexistent") + try: + fm.load_script(tutils.test_data.path("scripts/starterr.py")) + except ScriptException as e: + assert "ValueError" in str(e) assert len(fm.scripts) == 0 def test_getset_ignore(self): @@ -779,7 +783,7 @@ class TestFlowMaster: def test_script_reqerr(self): s = flow.State() fm = flow.FlowMaster(None, s) - assert not fm.load_script(tutils.test_data.path("scripts/reqerr.py")) + fm.load_script(tutils.test_data.path("scripts/reqerr.py")) f = tutils.tflow() fm.handle_clientconnect(f.client_conn) assert fm.handle_request(f) @@ -787,7 +791,7 @@ class TestFlowMaster: def test_script(self): s = flow.State() fm = flow.FlowMaster(None, s) - assert not fm.load_script(tutils.test_data.path("scripts/all.py")) + fm.load_script(tutils.test_data.path("scripts/all.py")) f = tutils.tflow(resp=True) fm.handle_clientconnect(f.client_conn) @@ -799,7 +803,7 @@ class TestFlowMaster: fm.handle_response(f) assert fm.scripts[0].ns["log"][-1] == "response" # load second script - assert not fm.load_script(tutils.test_data.path("scripts/all.py")) + fm.load_script(tutils.test_data.path("scripts/all.py")) assert len(fm.scripts) == 2 fm.handle_clientdisconnect(f.server_conn) assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" @@ -808,7 +812,7 @@ class TestFlowMaster: # unload first script fm.unload_scripts() assert len(fm.scripts) == 0 - assert not fm.load_script(tutils.test_data.path("scripts/all.py")) + fm.load_script(tutils.test_data.path("scripts/all.py")) f.error = tutils.terr() fm.handle_error(f) diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 8843ee62c..454736d4d 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -285,8 +285,7 @@ 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("scripts/stream_modify.py")) + self.master.load_script(tutils.test_data.path("scripts/stream_modify.py")) d = self.pathod('200:b"foo"') assert d.content == "bar" self.master.unload_scripts() @@ -511,8 +510,7 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): ssl = False def test_tcp_stream_modify(self): - self.master.load_script( - tutils.test_data.path("scripts/tcp_stream_modify.py")) + self.master.load_script(tutils.test_data.path("scripts/tcp_stream_modify.py")) self._tcpproxy_on() d = self.pathod('200:b"foo"')