Merge pull request #2258 from mhils/readfile

Integrate readstdin into readfile

10 files changed, 147 insertions, 152 deletions

diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py
index b4367d78..7a45106c 100644
--- a/mitmproxy/addons/__init__.py
+++ b/mitmproxy/addons/__init__.py
@@ -8,7 +8,6 @@ from mitmproxy.addons import disable_h2c
 from mitmproxy.addons import onboarding
 from mitmproxy.addons import proxyauth
 from mitmproxy.addons import replace
-from mitmproxy.addons import readfile
 from mitmproxy.addons import script
 from mitmproxy.addons import serverplayback
 from mitmproxy.addons import setheaders
@@ -38,6 +37,5 @@ def default_addons():
         stickycookie.StickyCookie(),
         streambodies.StreamBodies(),
         streamfile.StreamFile(),
-        readfile.ReadFile(),
         upstream_auth.UpstreamAuth(),
     ]
diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py
index 949da15d..05b6c309 100644
--- a/mitmproxy/addons/readfile.py
+++ b/mitmproxy/addons/readfile.py
@@ -1,38 +1,56 @@
 import os.path
+import sys
+import typing
 
 from mitmproxy import ctx
-from mitmproxy import io
 from mitmproxy import exceptions
+from mitmproxy import io
 
 
 class ReadFile:
     """
         An addon that handles reading from file on startup.
     """
-    def load_flows_file(self, path: str) -> int:
-        path = os.path.expanduser(path)
+
+    def load_flows(self, fo: typing.IO[bytes]) -> int:
         cnt = 0
+        freader = io.FlowReader(fo)
         try:
-            with open(path, "rb") as f:
-                freader = io.FlowReader(f)
-                for i in freader.stream():
-                    cnt += 1
-                    ctx.master.load_flow(i)
-                return cnt
-        except (IOError, exceptions.FlowReadException) as v:
+            for flow in freader.stream():
+                ctx.master.load_flow(flow)
+                cnt += 1
+        except (IOError, exceptions.FlowReadException) as e:
             if cnt:
-                ctx.log.warn(
-                    "Flow file corrupted - loaded %i flows." % cnt,
-                )
+                ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt)
             else:
                 ctx.log.error("Flow file corrupted.")
-            raise exceptions.FlowReadException(v)
+            raise exceptions.FlowReadException(str(e)) from e
+        else:
+            return cnt
+
+    def load_flows_from_path(self, path: str) -> int:
+        path = os.path.expanduser(path)
+        try:
+            with open(path, "rb") as f:
+                return self.load_flows(f)
+        except IOError as e:
+            ctx.log.error("Cannot load flows: {}".format(e))
+            raise exceptions.FlowReadException(str(e)) from e
 
     def running(self):
         if ctx.options.rfile:
             try:
-                self.load_flows_file(ctx.options.rfile)
-            except exceptions.FlowReadException as v:
-                raise exceptions.OptionsError(v)
+                self.load_flows_from_path(ctx.options.rfile)
+            except exceptions.FlowReadException as e:
+                raise exceptions.OptionsError(e) from e
             finally:
                 ctx.master.addons.trigger("processing_complete")
+
+
+class ReadFileStdin(ReadFile):
+    """Support the special case of "-" for reading from stdin"""
+    def load_flows_from_path(self, path: str) -> int:
+        if path == "-":
+            return self.load_flows(sys.stdin.buffer)
+        else:
+            return super().load_flows_from_path(path)
diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py
deleted file mode 100644
index 93a99f01..00000000
--- a/mitmproxy/addons/readstdin.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from mitmproxy import ctx
-from mitmproxy import io
-from mitmproxy import exceptions
-import sys
-
-
-class ReadStdin:
-    """
-        An addon that reads from stdin if we're not attached to (someting like)
-        a tty.
-    """
-    def running(self, stdin = sys.stdin):
-        if not stdin.isatty():
-            ctx.log.info("Reading from stdin")
-            try:
-                stdin.buffer.read(0)
-            except Exception as e:
-                ctx.log.warn("Cannot read from stdin: {}".format(e))
-                return
-            freader = io.FlowReader(stdin.buffer)
-            try:
-                for i in freader.stream():
-                    ctx.master.load_flow(i)
-            except exceptions.FlowReadException as e:
-                ctx.log.error("Error reading from stdin: %s" % e)
-            ctx.master.addons.trigger("processing_complete")
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index c1d584ac..e7a2c6ae 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -17,8 +17,9 @@ from mitmproxy import exceptions
 from mitmproxy import master
 from mitmproxy import io
 from mitmproxy import log
-from mitmproxy.addons import view
 from mitmproxy.addons import intercept
+from mitmproxy.addons import readfile
+from mitmproxy.addons import view
 from mitmproxy.tools.console import flowlist
 from mitmproxy.tools.console import flowview
 from mitmproxy.tools.console import grideditor
@@ -91,7 +92,12 @@ class ConsoleMaster(master.Master):
         signals.sig_add_log.connect(self.sig_add_log)
         self.addons.add(Logger())
         self.addons.add(*addons.default_addons())
-        self.addons.add(intercept.Intercept(), self.view, UnsupportedLog())
+        self.addons.add(
+            intercept.Intercept(),
+            self.view,
+            UnsupportedLog(),
+            readfile.ReadFile(),
+        )
 
         def sigint_handler(*args, **kwargs):
             self.prompt_for_exit()
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index 6329f6b7..4d0ccf4b 100644
--- a/mitmproxy/tools/dump.py
+++ b/mitmproxy/tools/dump.py
@@ -1,7 +1,7 @@
 from mitmproxy import addons
 from mitmproxy import options
 from mitmproxy import master
-from mitmproxy.addons import dumper, termlog, termstatus, readstdin, keepserving
+from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile
 
 
 class ErrorCheck:
@@ -30,7 +30,7 @@ class DumpMaster(master.Master):
         if with_dumper:
             self.addons.add(dumper.Dumper())
         self.addons.add(
-            readstdin.ReadStdin(),
             keepserving.KeepServing(),
+            readfile.ReadFileStdin(),
             self.errorcheck
         )
diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py
index 0db5a09f..c09fe0a2 100644
--- a/mitmproxy/tools/web/master.py
+++ b/mitmproxy/tools/web/master.py
@@ -7,6 +7,7 @@ from mitmproxy import log
 from mitmproxy import master
 from mitmproxy.addons import eventstore
 from mitmproxy.addons import intercept
+from mitmproxy.addons import readfile
 from mitmproxy.addons import termlog
 from mitmproxy.addons import view
 from mitmproxy.addons import termstatus
@@ -32,6 +33,7 @@ class WebMaster(master.Master):
         self.addons.add(*addons.default_addons())
         self.addons.add(
             intercept.Intercept(),
+            readfile.ReadFile(),
             self.view,
             self.events,
         )
diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py
index 5df4ea4b..628ea642 100644
--- a/mitmproxy/utils/typecheck.py
+++ b/mitmproxy/utils/typecheck.py
@@ -68,5 +68,7 @@ def check_type(name: str, value: typing.Any, typeinfo: typing.Any) -> None:
             return
         else:
             raise e
+    elif typename.startswith("typing.Any"):
+        return
     elif not isinstance(value, typeinfo):
         raise e
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py
index b30c147b..813aa10e 100644
--- a/test/mitmproxy/addons/test_readfile.py
+++ b/test/mitmproxy/addons/test_readfile.py
@@ -1,62 +1,104 @@
+import io
+from unittest import mock
+
+import pytest
+
+import mitmproxy.io
+from mitmproxy import exceptions
 from mitmproxy.addons import readfile
 from mitmproxy.test import taddons
 from mitmproxy.test import tflow
-from mitmproxy import io
-from mitmproxy import exceptions
-from unittest import mock
 
-import pytest
 
+@pytest.fixture
+def data():
+    f = io.BytesIO()
+
+    w = mitmproxy.io.FlowWriter(f)
+    flows = [
+        tflow.tflow(resp=True),
+        tflow.tflow(err=True),
+        tflow.ttcpflow(),
+        tflow.ttcpflow(err=True)
+    ]
+    for flow in flows:
+        w.add(flow)
 
-def write_data(path, corrupt=False):
-    with open(path, "wb") as tf:
-        w = io.FlowWriter(tf)
-        for i in range(3):
-            f = tflow.tflow(resp=True)
-            w.add(f)
-        for i in range(3):
-            f = tflow.tflow(err=True)
-            w.add(f)
-        f = tflow.ttcpflow()
-        w.add(f)
-        f = tflow.ttcpflow(err=True)
-        w.add(f)
-        if corrupt:
-            tf.write(b"flibble")
-
-
-@mock.patch('mitmproxy.master.Master.load_flow')
-def test_configure(mck, tmpdir):
-
-    rf = readfile.ReadFile()
-    with taddons.context() as tctx:
-        tf = str(tmpdir.join("tfile"))
-        write_data(tf)
-        tctx.configure(rf, rfile=str(tf))
-        assert not mck.called
-        rf.running()
-        assert mck.called
-
-        write_data(tf, corrupt=True)
-        tctx.configure(rf, rfile=str(tf))
-        with pytest.raises(exceptions.OptionsError):
+    f.seek(0)
+    return f
+
+
+@pytest.fixture
+def corrupt_data():
+    f = data()
+    f.seek(0, io.SEEK_END)
+    f.write(b"qibble")
+    f.seek(0)
+    return f
+
+
+class TestReadFile:
+    @mock.patch('mitmproxy.master.Master.load_flow')
+    def test_configure(self, mck, tmpdir, data, corrupt_data):
+        rf = readfile.ReadFile()
+        with taddons.context() as tctx:
+            tf = tmpdir.join("tfile")
+
+            tf.write(data.getvalue())
+            tctx.configure(rf, rfile=str(tf))
+            assert not mck.called
             rf.running()
+            assert mck.called
 
+            tf.write(corrupt_data.getvalue())
+            tctx.configure(rf, rfile=str(tf))
+            with pytest.raises(exceptions.OptionsError):
+                rf.running()
 
-@mock.patch('mitmproxy.master.Master.load_flow')
-def test_corruption(mck, tmpdir):
+    @mock.patch('mitmproxy.master.Master.load_flow')
+    def test_corrupt(self, mck, corrupt_data):
+        rf = readfile.ReadFile()
+        with taddons.context() as tctx:
+            with pytest.raises(exceptions.FlowReadException):
+                rf.load_flows(io.BytesIO(b"qibble"))
+            assert not mck.called
+            assert len(tctx.master.logs) == 1
 
-    rf = readfile.ReadFile()
-    with taddons.context() as tctx:
-        with pytest.raises(exceptions.FlowReadException):
-            rf.load_flows_file("nonexistent")
-        assert not mck.called
-        assert len(tctx.master.logs) == 1
+            with pytest.raises(exceptions.FlowReadException):
+                rf.load_flows(corrupt_data)
+            assert mck.called
+            assert len(tctx.master.logs) == 2
+
+    def test_nonexisting_file(self):
+        rf = readfile.ReadFile()
+        with taddons.context() as tctx:
+            with pytest.raises(exceptions.FlowReadException):
+                rf.load_flows_from_path("nonexistent")
+            assert len(tctx.master.logs) == 1
+
+
+class TestReadFileStdin:
+    @mock.patch('mitmproxy.master.Master.load_flow')
+    @mock.patch('sys.stdin')
+    def test_stdin(self, stdin, load_flow, data, corrupt_data):
+        rf = readfile.ReadFileStdin()
+        with taddons.context() as tctx:
+            stdin.buffer = data
+            tctx.configure(rf, rfile="-")
+            assert not load_flow.called
+            rf.running()
+            assert load_flow.called
 
-        tfc = str(tmpdir.join("tfile"))
-        write_data(tfc, corrupt=True)
+            stdin.buffer = corrupt_data
+            tctx.configure(rf, rfile="-")
+            with pytest.raises(exceptions.OptionsError):
+                rf.running()
 
-        with pytest.raises(exceptions.FlowReadException):
-            rf.load_flows_file(tfc)
-        assert mck.called
-        assert len(tctx.master.logs) == 2
+    @mock.patch('mitmproxy.master.Master.load_flow')
+    def test_normal(self, load_flow, tmpdir, data):
+        rf = readfile.ReadFileStdin()
+        with taddons.context():
+            tfile = tmpdir.join("tfile")
+            tfile.write(data.getvalue())
+            rf.load_flows_from_path(str(tfile))
+            assert load_flow.called
diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py
deleted file mode 100644
index 76b01f4f..00000000
--- a/test/mitmproxy/addons/test_readstdin.py
+++ /dev/null
@@ -1,53 +0,0 @@
-
-import io
-from mitmproxy.addons import readstdin
-from mitmproxy.test import taddons
-from mitmproxy.test import tflow
-import mitmproxy.io
-from unittest import mock
-
-
-def gen_data(corrupt=False):
-    tf = io.BytesIO()
-    w = mitmproxy.io.FlowWriter(tf)
-    for i in range(3):
-        f = tflow.tflow(resp=True)
-        w.add(f)
-    for i in range(3):
-        f = tflow.tflow(err=True)
-        w.add(f)
-    f = tflow.ttcpflow()
-    w.add(f)
-    f = tflow.ttcpflow(err=True)
-    w.add(f)
-    if corrupt:
-        tf.write(b"flibble")
-    tf.seek(0)
-    return tf
-
-
-class mStdin:
-    def __init__(self, d):
-        self.buffer = d
-
-    def isatty(self):
-        return False
-
-
-@mock.patch('mitmproxy.master.Master.load_flow')
-def test_read(m, tmpdir):
-    rf = readstdin.ReadStdin()
-    with taddons.context() as tctx:
-        assert not m.called
-        rf.running(stdin=mStdin(gen_data()))
-        assert m.called
-
-        rf.running(stdin=mStdin(None))
-        assert tctx.master.logs
-        tctx.master.clear()
-
-        m.reset_mock()
-        assert not m.called
-        rf.running(stdin=mStdin(gen_data(corrupt=True)))
-        assert m.called
-        assert tctx.master.logs
diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py
index d99a914f..fd0c6e0c 100644
--- a/test/mitmproxy/utils/test_typecheck.py
+++ b/test/mitmproxy/utils/test_typecheck.py
@@ -79,3 +79,9 @@ def test_check_io():
     typecheck.check_type("foo", io.StringIO(), typing.IO[str])
     with pytest.raises(TypeError):
         typecheck.check_type("foo", "foo", typing.IO[str])
+
+
+def test_check_any():
+    typecheck.check_type("foo", 42, typing.Any)
+    typecheck.check_type("foo", object(), typing.Any)
+    typecheck.check_type("foo", None, typing.Any)