Merge pull request #2971 from kira0204/user-script-exception

A unified handler for user's scripts exception

2 files changed, 38 insertions, 6 deletions

diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index 6f0d1e28..f5c4d599 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -5,6 +5,7 @@ import time
 import sys
 import types
 import typing
+import traceback
 
 from mitmproxy import addonmanager
 from mitmproxy import exceptions
@@ -36,6 +37,25 @@ def load_script(path: str) -> types.ModuleType:
         sys.path[:] = oldpath
 
 
+def script_error_handler(path, exc, msg="", tb=False):
+    """
+        Handles all the user's script errors with
+        an optional traceback
+    """
+    exception = type(exc).__name__
+    if msg:
+        exception = msg
+    lineno = ""
+    if hasattr(exc, "lineno"):
+        lineno = str(exc.lineno)
+    log_msg = "in Script {}:{} {}".format(path, lineno, exception)
+    if tb:
+        etype, value, tback = sys.exc_info()
+        tback = addonmanager.cut_traceback(tback, "invoke_addon")
+        log_msg = log_msg.join(["\n"] + traceback.format_exception(etype, value, tback))
+    ctx.log.error(log_msg)
+
+
 class Script:
     """
         An addon that manages a single script.
@@ -53,7 +73,7 @@ class Script:
         self.last_load = 0
         self.last_mtime = 0
         if not os.path.isfile(self.fullpath):
-            raise exceptions.OptionsError('No such script: "%s"' % self.fullpath)
+            raise exceptions.OptionsError('No such script')
 
     @property
     def addons(self):
@@ -128,13 +148,13 @@ class ScriptLoader:
                 for evt, arg in eventsequence.iterate(f):
                     ctx.master.addons.invoke_addon(s, evt, arg)
         except exceptions.OptionsError as e:
-            raise exceptions.CommandError("Error running script: %s" % e) from e
+            script_error_handler(path, e, msg=str(e))
 
     def configure(self, updated):
         if "scripts" in updated:
             for s in ctx.options.scripts:
                 if ctx.options.scripts.count(s) > 1:
-                    raise exceptions.OptionsError("Duplicate script: %s" % s)
+                    raise exceptions.OptionsError("Duplicate script")
 
             for a in self.addons[:]:
                 if a.path not in ctx.options.scripts:
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index dc21e6fd..79fa22dc 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -183,9 +183,9 @@ class TestScriptLoader:
 
     def test_script_run_nonexistent(self):
         sc = script.ScriptLoader()
-        with taddons.context(sc):
-            with pytest.raises(exceptions.CommandError):
-                sc.script_run([tflow.tflow(resp=True)], "/")
+        with taddons.context(sc) as tctx:
+            sc.script_run([tflow.tflow(resp=True)], "/")
+            tctx.master.has_log("/: No such script")
 
     def test_simple(self):
         sc = script.ScriptLoader()
@@ -243,6 +243,18 @@ class TestScriptLoader:
             tctx.invoke(sc, "tick")
             assert len(tctx.master.addons) == 1
 
+    def test_script_error_handler(self):
+        path = "/sample/path/example.py"
+        exc = SyntaxError
+        msg = "Error raised"
+        tb = True
+        with taddons.context() as tctx:
+            script.script_error_handler(path, exc, msg, tb)
+            assert tctx.master.has_log("/sample/path/example.py")
+            assert tctx.master.has_log("Error raised")
+            assert tctx.master.has_log("lineno")
+            assert tctx.master.has_log("NoneType")
+
     def test_order(self):
         rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder")
         sc = script.ScriptLoader()