From a3329c5a80b386a12575002d8324eb783428ba16 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 22 Sep 2016 02:03:12 -0700
Subject: scripts: cut traceback properly, fix #1467

---
 mitmproxy/builtins/script.py           | 33 ++++++++++++++++++++++++++++-----
 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()
-- 
cgit v1.2.3