Extract flow reading into addons

This patch moves the final pieces of master functionality into addons.

- Add a ReadFile addon to read from file
- Add a separate ReadStdin addon to read from stdin, only used by mitmdump
- Remove all methods that know about io and serialization from master.Master

8 files changed, 95 insertions, 70 deletions

diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py
index 7a45106c..b4367d78 100644
--- a/mitmproxy/addons/__init__.py
+++ b/mitmproxy/addons/__init__.py
@@ -8,6 +8,7 @@ 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
@@ -37,5 +38,6 @@ 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
new file mode 100644
index 00000000..a4b92444
--- /dev/null
+++ b/mitmproxy/addons/readfile.py
@@ -0,0 +1,50 @@
+import os.path
+
+from mitmproxy import ctx
+from mitmproxy import io
+from mitmproxy import exceptions
+
+
+class ReadFile:
+    """
+        An addon that handles reading from file on startup.
+    """
+    def __init__(self):
+        self.path = None
+        self.keepserving = False
+
+    def load_flows_file(self, path: str) -> int:
+        path = os.path.expanduser(path)
+        cnt = 0
+        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:
+            if cnt:
+                ctx.log.warn(
+                    "Flow file corrupted - loaded %i flows." % cnt,
+                )
+            else:
+                ctx.log.error("Flow file corrupted.")
+            raise exceptions.FlowReadException(v)
+
+    def configure(self, options, updated):
+        if "keepserving" in updated:
+            self.keepserving = options.keepserving
+        if "rfile" in updated and options.rfile:
+            self.path = options.rfile
+
+    def running(self):
+        if self.path:
+            try:
+                self.load_flows_file(self.path)
+            except exceptions.FlowReadException as v:
+                raise exceptions.OptionsError(v)
+            finally:
+                self.path = None
+                if not self.keepserving:
+                    ctx.master.shutdown()
diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py
new file mode 100644
index 00000000..e45d25b8
--- /dev/null
+++ b/mitmproxy/addons/readstdin.py
@@ -0,0 +1,34 @@
+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 __init__(self):
+        self.keepserving = False
+
+    def configure(self, options, updated):
+        if "keepserving" in updated:
+            self.keepserving = options.keepserving
+
+    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)
+            if not self.keepserving:
+                ctx.master.shutdown()
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 79747a97..69359de6 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -1,8 +1,6 @@
-import os
 import threading
 import contextlib
 import queue
-import sys
 
 from mitmproxy import addonmanager
 from mitmproxy import options
@@ -12,7 +10,6 @@ from mitmproxy import exceptions
 from mitmproxy import connections
 from mitmproxy import http
 from mitmproxy import log
-from mitmproxy import io
 from mitmproxy.proxy.protocol import http_replay
 from mitmproxy.types import basethread
 import mitmproxy.net.http
@@ -160,33 +157,6 @@ class Master:
         for e, o in eventsequence.iterate(f):
             getattr(self, e)(o)
 
-    def load_flows(self, fr: io.FlowReader) -> int:
-        """
-            Load flows from a FlowReader object.
-        """
-        cnt = 0
-        for i in fr.stream():
-            cnt += 1
-            self.load_flow(i)
-        return cnt
-
-    def load_flows_file(self, path: str) -> int:
-        path = os.path.expanduser(path)
-        try:
-            if path == "-":
-                try:
-                    sys.stdin.buffer.read(0)
-                except Exception as e:
-                    raise IOError("Cannot read from stdin: {}".format(e))
-                freader = io.FlowReader(sys.stdin.buffer)
-                return self.load_flows(freader)
-            else:
-                with open(path, "rb") as f:
-                    freader = io.FlowReader(f)
-                    return self.load_flows(freader)
-        except IOError as v:
-            raise exceptions.FlowReadException(v.strerror)
-
     def replay_request(
             self,
             f: http.HTTPFlow,
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index e75105cf..b6339817 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -256,19 +256,6 @@ class ConsoleMaster(master.Master):
         )
         self.ab = statusbar.ActionBar()
 
-        if self.options.rfile:
-            ret = self.load_flows_path(self.options.rfile)
-            if ret and self.view.store_count():
-                signals.add_log(
-                    "File truncated or corrupted. "
-                    "Loaded as many flows as possible.",
-                    "error"
-                )
-            elif ret and not self.view.store_count():
-                self.shutdown()
-                print("Could not load file: {}".format(ret), file=sys.stderr)
-                sys.exit(1)
-
         self.loop.set_alarm_in(0.01, self.ticker)
 
         self.loop.set_alarm_in(
@@ -289,7 +276,10 @@ class ConsoleMaster(master.Master):
             print("Shutting down...", file=sys.stderr)
         finally:
             sys.stderr.flush()
-            self.shutdown()
+            super().shutdown()
+
+    def shutdown(self):
+        raise urwid.ExitMainLoop
 
     def view_help(self, helpctx):
         signals.push_view_state.send(
@@ -402,7 +392,7 @@ class ConsoleMaster(master.Master):
 
     def quit(self, a):
         if a != "n":
-            raise urwid.ExitMainLoop
+            self.shutdown()
 
     def clear_events(self):
         self.logbuffer[:] = []
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index 4bfe2dc4..be83fb1d 100644
--- a/mitmproxy/tools/dump.py
+++ b/mitmproxy/tools/dump.py
@@ -1,9 +1,8 @@
 from mitmproxy import controller
-from mitmproxy import exceptions
 from mitmproxy import addons
 from mitmproxy import options
 from mitmproxy import master
-from mitmproxy.addons import dumper, termlog, termstatus
+from mitmproxy.addons import dumper, termlog, termstatus, readstdin
 
 
 class DumpMaster(master.Master):
@@ -22,21 +21,9 @@ class DumpMaster(master.Master):
         self.addons.add(*addons.default_addons())
         if with_dumper:
             self.addons.add(dumper.Dumper())
-
-        if options.rfile:
-            try:
-                self.load_flows_file(options.rfile)
-            except exceptions.FlowReadException as v:
-                self.add_log("Flow file corrupted.", "error")
-                raise exceptions.OptionsError(v)
+        self.addons.add(readstdin.ReadStdin())
 
     @controller.handler
     def log(self, e):
         if e.level == "error":
             self.has_errored = True
-
-    def run(self):  # pragma: no cover
-        if self.options.rfile and not self.options.keepserving:
-            self.addons.done()
-            return
-        super().run()
diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py
index eddaa3e1..002513b9 100644
--- a/mitmproxy/tools/web/app.py
+++ b/mitmproxy/tools/web/app.py
@@ -230,7 +230,8 @@ class DumpFlows(RequestHandler):
     def post(self):
         self.view.clear()
         bio = BytesIO(self.filecontents)
-        self.master.load_flows(io.FlowReader(bio))
+        for i in io.FlowReader(bio).stream():
+            self.master.load_flow(i)
         bio.close()
 
 
diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py
index 8c7f579d..e28bd002 100644
--- a/mitmproxy/tools/web/master.py
+++ b/mitmproxy/tools/web/master.py
@@ -3,7 +3,6 @@ import webbrowser
 import tornado.httpserver
 import tornado.ioloop
 from mitmproxy import addons
-from mitmproxy import exceptions
 from mitmproxy import log
 from mitmproxy import master
 from mitmproxy.addons import eventstore
@@ -42,14 +41,6 @@ class WebMaster(master.Master):
         )
         # This line is just for type hinting
         self.options = self.options  # type: Options
-        if options.rfile:
-            try:
-                self.load_flows_file(options.rfile)
-            except exceptions.FlowReadException as v:
-                self.add_log(
-                    "Could not read flow file: %s" % v,
-                    "error"
-                )
 
     def _sig_view_add(self, view, flow):
         app.ClientConnection.broadcast(