improve script loading

9 files changed, 76 insertions, 54 deletions

diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py
index 5b980572..e75aed86 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 597a7e7a..46ff348e 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 d224f1aa..0e0ccc62 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 86bf75ae..52fd36d1 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 9292e76a..c63c1efa 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 4d0f73fa..484025b4 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 b560d9a1..c401a6b9 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 145e91cf..84a25689 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 8843ee62..454736d4 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"')