scripts: redirect stdout to ctx.log.warn

Redirect messages written to stdout in scripts to ctx.log.warn.
(closes #1530)

Signed-off-by: Nikos Filippakis <nikolaos.filippakis@cern.ch>

mitmproxy/addons/script.py

@@ -65,13 +65,27 @@ def cut_traceback(tb, func_name):
         return tb
 
 
+class StreamLog:
+    """
+        A class for redirecting output using contextlib.
+    """
+    def __init__(self, log):
+        self.log = log
+
+    def write(self, buf):
+        if buf.strip():
+            self.log(buf)
+
+
 @contextlib.contextmanager
 def scriptenv(path, args):
     oldargs = sys.argv
     sys.argv = [path] + args
     script_dir = os.path.dirname(os.path.abspath(path))
     sys.path.append(script_dir)
+    stdout_replacement = StreamLog(ctx.log.warn)
     try:
+        with contextlib.redirect_stdout(stdout_replacement):
             yield
     except SystemExit as v:
         ctx.log.error("Script exited with code %s" % v.code)

test/mitmproxy/addons/test_script.py

@@ -5,6 +5,7 @@ import re
 import watchdog.events
 import pytest
 
+from unittest import mock
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 from mitmproxy.test import taddons
@@ -97,6 +98,7 @@ class TestParseCommand:
 
 
 def test_load_script():
+    with taddons.context():
         ns = script.load_script(
             tutils.test_data.path(
                 "mitmproxy/data/addonscripts/recorder.py"
@@ -105,6 +107,19 @@ def test_load_script():
         assert ns.start
 
 
+def test_script_print_stdout():
+    with taddons.context() as tctx:
+        with mock.patch('mitmproxy.ctx.log.warn') as mock_warn:
+            with script.scriptenv("path", []):
+                ns = script.load_script(
+                    tutils.test_data.path(
+                        "mitmproxy/data/addonscripts/print.py"
+                    ), []
+                )
+                ns.start(tctx.options)
+        mock_warn.assert_called_once_with("stdoutprint")
+
+
 class TestScript:
     def test_simple(self):
         with taddons.context():

test/mitmproxy/data/addonscripts/print.py

@@ -0,0 +1,2 @@
+def start(opts):
+    print("stdoutprint")