Script cleanup: editing in console, Python3 compatibility fixes

7 files changed, 38 insertions, 61 deletions

diff --git a/mitmproxy/console/grideditor.py b/mitmproxy/console/grideditor.py
index 9fa51ccb..f304de57 100644
--- a/mitmproxy/console/grideditor.py
+++ b/mitmproxy/console/grideditor.py
@@ -6,11 +6,12 @@ import re
 
 import urwid
 
+from mitmproxy import exceptions
 from mitmproxy import filt
-from mitmproxy import script
-from mitmproxy import utils
+from mitmproxy.builtins import script
 from mitmproxy.console import common
 from mitmproxy.console import signals
+from netlib import strutils
 from netlib.http import cookies
 from netlib.http import user_agents
 
@@ -55,7 +56,7 @@ class TextColumn:
             o = editor.walker.get_current_value()
             if o is not None:
                 n = editor.master.spawn_editor(o.encode("string-escape"))
-                n = utils.clean_hanging_newline(n)
+                n = strutils.clean_hanging_newline(n)
                 editor.walker.set_current_value(n, False)
                 editor.walker._modified()
         elif key in ["enter"]:
@@ -643,8 +644,8 @@ class ScriptEditor(GridEditor):
 
     def is_error(self, col, val):
         try:
-            script.Script.parse_command(val)
-        except script.ScriptException as e:
+            script.parse_command(val)
+        except exceptions.AddonError as e:
             return str(e)
 
 
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index bc373a2b..64bd9f0a 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -248,23 +248,6 @@ class ConsoleMaster(flow.FlowMaster):
         if options.server_replay:
             self.server_playback_path(options.server_replay)
 
-        if options.scripts:
-            for i in options.scripts:
-                try:
-                    self.load_script(i)
-                except exceptions.ScriptException as e:
-                    print("Script load error: {}".format(e), file=sys.stderr)
-                    sys.exit(1)
-
-        if options.outfile:
-            err = self.start_stream_to_path(
-                options.outfile[0],
-                options.outfile[1]
-            )
-            if err:
-                print("Stream file error: {}".format(err), file=sys.stderr)
-                sys.exit(1)
-
         self.view_stack = []
 
         if options.app:
@@ -685,20 +668,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.refresh_focus()
 
     def edit_scripts(self, scripts):
-        commands = [x[0] for x in scripts]  # remove outer array
-        if commands == [s.command for s in self.scripts]:
-            return
-
-        self.unload_scripts()
-        for command in commands:
-            try:
-                self.load_script(command)
-            except exceptions.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)
+        self.options.scripts = [x[0] for x in scripts]
 
     def stop_client_playback_prompt(self, a):
         if a != "n":
diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py
index d363ba74..d8824b05 100644
--- a/mitmproxy/console/options.py
+++ b/mitmproxy/console/options.py
@@ -54,7 +54,7 @@ class Options(urwid.WidgetWrap):
                 select.Option(
                     "Scripts",
                     "S",
-                    lambda: master.scripts,
+                    lambda: master.options.scripts,
                     self.scripts
                 ),
 
@@ -160,12 +160,14 @@ class Options(urwid.WidgetWrap):
         self.master.replacehooks.clear()
         self.master.set_ignore_filter([])
         self.master.set_tcp_filter([])
-        self.master.scripts = []
 
-        self.master.options.anticache = False
-        self.master.options.anticomp = False
-        self.master.options.stickyauth = None
-        self.master.options.stickycookie = None
+        self.master.options.update(
+            scripts = [],
+            anticache = False,
+            anticomp = False,
+            stickyauth = None,
+            stickycookie = None
+        )
 
         self.master.state.default_body_view = contentviews.get("Auto")
 
@@ -234,7 +236,7 @@ class Options(urwid.WidgetWrap):
         self.master.view_grideditor(
             grideditor.ScriptEditor(
                 self.master,
-                [[i.command] for i in self.master.scripts],
+                [[i] for i in self.master.options.scripts],
                 self.master.edit_scripts
             )
         )
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index fc41869c..e7a700a6 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -218,14 +218,13 @@ class StatusBar(urwid.WidgetWrap):
                 dst.address.host,
                 dst.address.port
             ))
-        if self.master.scripts:
+        if self.master.options.scripts:
             r.append("[")
             r.append(("heading_key", "s"))
-            r.append("cripts:%s]" % len(self.master.scripts))
-        # r.append("[lt:%0.3f]"%self.master.looptime)
+            r.append("cripts:%s]" % len(self.master.options.scripts))
 
-        if self.master.stream:
-            r.append("[W:%s]" % self.master.stream_path)
+        if self.master.options.outfile:
+            r.append("[W:%s]" % self.master.outfile[0])
 
         return r
 
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index d5243bcb..92d1153b 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -1,11 +1,10 @@
-from test.mitmproxy import tutils
+from test.mitmproxy import tutils, mastertest
 from mitmproxy import controller
 from mitmproxy.builtins import script
 from mitmproxy import options
 from mitmproxy.flow import master
 from mitmproxy.flow import state
 import time
-from .. import mastertest, tutils
 
 
 class Thing:
diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index 9c8edb29..ef97219c 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -1,5 +1,6 @@
 import json
 
+import sys
 import os.path
 from mitmproxy.flow import master
 from mitmproxy.flow import state
@@ -42,19 +43,19 @@ class TestScripts(mastertest.MasterTest):
     def test_custom_contentviews(self):
         m, sc = tscript("custom_contentviews.py")
         pig = contentviews.get("pig_latin_HTML")
-        _, fmt = pig("<html>test!</html>")
-        assert any('esttay!' in val[0][1] for val in fmt)
-        assert not pig("gobbledygook")
+        _, fmt = pig(b"<html>test!</html>")
+        assert any(b'esttay!' in val[0][1] for val in fmt)
+        assert not pig(b"gobbledygook")
 
     def test_iframe_injector(self):
         with tutils.raises(ScriptError):
             tscript("iframe_injector.py")
 
         m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe")
-        flow = tutils.tflow(resp=netutils.tresp(content="<html>mitmproxy</html>"))
+        flow = tutils.tflow(resp=netutils.tresp(content=b"<html>mitmproxy</html>"))
         self.invoke(m, "response", flow)
         content = flow.response.content
-        assert 'iframe' in content and 'evil_iframe' in content
+        assert b'iframe' in content and b'evil_iframe' in content
 
     def test_modify_form(self):
         m, sc = tscript("modify_form.py")
@@ -63,11 +64,11 @@ class TestScripts(mastertest.MasterTest):
         f = tutils.tflow(req=netutils.treq(headers=form_header))
         self.invoke(m, "request", f)
 
-        assert f.request.urlencoded_form["mitmproxy"] == "rocks"
+        assert f.request.urlencoded_form[b"mitmproxy"] == b"rocks"
 
         f.request.headers["content-type"] = ""
         self.invoke(m, "request", f)
-        assert list(f.request.urlencoded_form.items()) == [("foo", "bar")]
+        assert list(f.request.urlencoded_form.items()) == [(b"foo", b"bar")]
 
     def test_modify_querystring(self):
         m, sc = tscript("modify_querystring.py")
@@ -85,9 +86,9 @@ class TestScripts(mastertest.MasterTest):
             tscript("modify_response_body.py")
 
         m, sc = tscript("modify_response_body.py", "mitmproxy rocks")
-        f = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy"))
+        f = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy"))
         self.invoke(m, "response", f)
-        assert f.response.content == "I <3 rocks"
+        assert f.response.content == b"I <3 rocks"
 
     def test_redirect_requests(self):
         m, sc = tscript("redirect_requests.py")
@@ -96,6 +97,11 @@ class TestScripts(mastertest.MasterTest):
         assert f.request.host == "mitmproxy.org"
 
     def test_har_extractor(self):
+        if sys.version_info >= (3, 0):
+            with tutils.raises("does not work on Python 3"):
+                tscript("har_extractor.py")
+            return
+
         with tutils.raises(ScriptError):
             tscript("har_extractor.py")
 
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index b1ca6910..a5196dae 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -293,7 +293,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
         )
         self.master.addons.add(s)
         d = self.pathod('200:b"foo"')
-        assert d.content == "bar"
+        assert d.content == b"bar"
         self.master.addons.remove(s)
 
 
@@ -523,7 +523,7 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
         self._tcpproxy_on()
         d = self.pathod('200:b"foo"')
         self._tcpproxy_off()
-        assert d.content == "bar"
+        assert d.content == b"bar"
         self.master.addons.remove(s)