diff options
| -rw-r--r-- | mitmproxy/builtins/script.py | 33 | ||||
| -rw-r--r-- | test/mitmproxy/builtins/test_script.py | 30 |
2 files changed, 55 insertions, 8 deletions
diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index 1ebec873..ba62d416 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -53,6 +53,33 @@ def parse_command(command): return args[0], args[1:] +def cut_traceback(tb, func_name): + """ + Cut off a traceback at the function with the given name. + The func_name's frame is excluded. + + Args: + tb: traceback object, as returned by sys.exc_info()[2] + func_name: function name + + Returns: + Reduced traceback. + """ + tb_orig = tb + + for _, _, fname, _ in traceback.extract_tb(tb): + tb = tb.tb_next + if fname == func_name: + break + + if tb is None: + # We could not find the method, take the full stack trace. + # This may happen on some Python interpreters/flavors (e.g. PyInstaller). + return tb_orig + else: + return tb + + @contextlib.contextmanager def scriptenv(path, args): oldargs = sys.argv @@ -63,11 +90,7 @@ def scriptenv(path, args): yield except Exception: etype, value, 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 + tb = cut_traceback(tb, "scriptenv").tb_next ctx.log.error( "Script error: %s" % "".join( traceback.format_exception(etype, value, tb) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 09e5bc92..24337760 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -1,10 +1,12 @@ -import time +import traceback -from mitmproxy.builtins import script +import sys +import time from mitmproxy import exceptions +from mitmproxy import options +from mitmproxy.builtins import script from mitmproxy.flow import master from mitmproxy.flow import state -from mitmproxy import options from .. import tutils, mastertest @@ -104,6 +106,10 @@ class TestScript(mastertest.MasterTest): f = tutils.tflow(resp=True) m.request(f) assert m.event_log[0][0] == "error" + assert len(m.event_log[0][1].splitlines()) == 6 + assert 'addonscripts/error.py", line 7, in request' in m.event_log[0][1] + assert 'addonscripts/error.py", line 3, in mkerr' in m.event_log[0][1] + assert m.event_log[0][1].endswith("ValueError: Error!\n") def test_duplicate_flow(self): s = state.State() @@ -136,6 +142,24 @@ class TestScript(mastertest.MasterTest): ] +class TestCutTraceback: + def raise_(self, i): + if i > 0: + self.raise_(i - 1) + raise RuntimeError() + + def test_simple(self): + try: + self.raise_(4) + except RuntimeError: + tb = sys.exc_info()[2] + tb_cut = script.cut_traceback(tb, "test_simple") + assert len(traceback.extract_tb(tb_cut)) == 5 + + tb_cut2 = script.cut_traceback(tb, "nonexistent") + assert len(traceback.extract_tb(tb_cut2)) == len(traceback.extract_tb(tb)) + + class TestScriptLoader(mastertest.MasterTest): def test_run_once(self): s = state.State() |