From 0ac3227b7bb27a76b6f49aa1605f1b1887a01a10 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 5 Feb 2015 14:44:45 +0100
Subject: [PATCH] clean up flow reading

---
 libmproxy/console/__init__.py | 47 ++++++++++++++++-------------------
 libmproxy/dump.py             | 15 ++++++-----
 libmproxy/flow.py             | 19 ++++++++++++++
 test/test_dump.py             |  2 +-
 4 files changed, 48 insertions(+), 35 deletions(-)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 6aabf5ed5..0db068323 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -499,38 +499,33 @@ class ConsoleMaster(flow.FlowMaster):
         self.view_flowlist()
 
     def _readflow(self, paths):
+        """
+        Utitility function that reads a list of flows
+        or prints an error to the UI if that fails.
+        Returns
+            - None, if there was an error.
+            - a list of flows, otherwise.
+        """
         try:
-            flows = []
-            for path in paths:
-                path = os.path.expanduser(path)
-                with file(path, "rb") as f:
-                    flows.extend(list(flow.FlowReader(f).stream()))
-        except (IOError, flow.FlowReadError), v:
-            return True, v.strerror
-        return False, flows
+            return flow.read_flows_from_paths(paths)
+        except flow.FlowReadError as e:
+            if not self.statusbar:
+                print >> sys.stderr, e.strerror
+                sys.exit(1)
+            else:
+                self.statusbar.message(e.strerror)
+            return None
 
     def client_playback_path(self, path):
-        err, ret = self._readflow(path)
-        if err:
-            if not self.statusbar:
-                print >> sys.stderr, ret
-                sys.exit(1)
-            else:
-                self.statusbar.message(ret)
-        else:
-            self.start_client_playback(ret, False)
+        flows = self._readflow(path)
+        if flows:
+            self.start_client_playback(flows, False)
 
     def server_playback_path(self, path):
-        err, ret = self._readflow(path)
-        if err:
-            if not self.statusbar:
-                print >> sys.stderr, ret
-                sys.exit(1)
-            else:
-                self.statusbar.message(ret)
-        else:
+        flows = self._readflow(path)
+        if flows:
             self.start_server_playback(
-                ret,
+                flows,
                 self.killextra, self.rheaders,
                 False, self.nopop,
                 self.options.replay_ignore_params, self.options.replay_ignore_content, self.options.replay_ignore_payload_params
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 70f175f6a..59ddcf5f0 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -144,15 +144,14 @@ class DumpMaster(flow.FlowMaster):
             self.start_app(self.o.app_host, self.o.app_port)
 
     def _readflow(self, paths):
+        """
+        Utitility function that reads a list of flows
+        or raises a DumpError if that fails.
+        """
         try:
-            flows = []
-            for path in paths:
-                path = os.path.expanduser(path)
-                with file(path, "rb") as f:
-                    flows.extend(list(flow.FlowReader(f).stream()))
-        except (IOError, flow.FlowReadError), v:
-            raise DumpError(v.strerror)
-        return flows
+            return flow.read_flows_from_paths(paths)
+        except flow.FlowReadError as e:
+            raise DumpError(e.strerror)
 
     def add_event(self, e, level="info"):
         needed = dict(error=0, info=1, debug=2).get(level, 1)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 97ebc572b..f9e2b94d4 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -945,6 +945,25 @@ class FlowMaster(controller.Master):
         self.stream = None
 
 
+def read_flows_from_paths(paths):
+    """
+    Given a list of filepaths, read all flows and return a list of them.
+    From a performance perspective, streaming would be advisable -
+    however, if there's an error with one of the files, we want it to be raised immediately.
+
+    If an error occurs, a FlowReadError will be raised.
+    """
+    try:
+        flows = []
+        for path in paths:
+            path = os.path.expanduser(path)
+            with file(path, "rb") as f:
+                flows.extend(FlowReader(f).stream())
+    except IOError as e:
+        raise FlowReadError(e.strerror)
+    return flows
+
+
 class FlowWriter:
     def __init__(self, fo):
         self.fo = fo
diff --git a/test/test_dump.py b/test/test_dump.py
index 69df912f5..11a024e38 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -75,7 +75,7 @@ class TestDumpMaster:
     def test_replay(self):
         cs = StringIO()
 
-        o = dump.Options(server_replay="nonexistent", kill=True)
+        o = dump.Options(server_replay=["nonexistent"], kill=True)
         tutils.raises(dump.DumpError, dump.DumpMaster, None, o, outfile=cs)
 
         with tutils.tmpdir() as t: