commands: view.load

Plus replace the flow list keybinding.

7 files changed, 39 insertions, 27 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 794e7617..a629ceb9 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -20,6 +20,7 @@ from mitmproxy import flowfilter
 from mitmproxy import exceptions
 from mitmproxy import command
 from mitmproxy import ctx
+from mitmproxy import io
 from mitmproxy import http  # noqa
 
 # The underlying sorted list implementation expects the sort key to be stable
@@ -265,6 +266,17 @@ class View(collections.Sequence):
         """
         return self._store.get(flow_id)
 
+    @command.command("view.load")
+    def load_file(self, path: str) -> None:
+        """
+            Load flows into the view, without processing them with addons.
+        """
+        for i in io.FlowReader(open(path, "rb")).stream():
+            # Do this to get a new ID, so we can load the same file N times and
+            # get new flows each time. It would be more efficient to just have a
+            # .newid() method or something.
+            self.add([i.copy()])
+
     @command.command("view.go")
     def go(self, dst: int) -> None:
         """
diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py
index d4460273..812ca7a8 100644
--- a/mitmproxy/tools/console/common.py
+++ b/mitmproxy/tools/console/common.py
@@ -395,13 +395,16 @@ def raw_format_flow(f, flow):
 
 
 def format_flow(f, focus, extended=False, hostheader=False, max_url_len=False):
+    acked = False
+    if f.reply and f.reply.state == "committed":
+        acked = True
     d = dict(
         focus=focus,
         extended=extended,
         max_url_len=max_url_len,
 
         intercepted = f.intercepted,
-        acked = f.reply.state == "committed",
+        acked = acked,
 
         req_timestamp = f.request.timestamp_start,
         req_is_replay = f.request.is_replay,
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 7364524f..4ffed15f 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -227,13 +227,7 @@ class FlowListBox(urwid.ListBox):
 
     def keypress(self, size, key):
         key = common.shortcuts(key)
-        if key == "L":
-            signals.status_prompt_path.send(
-                self,
-                prompt = "Load flows",
-                callback = self.master.load_flows_callback
-            )
-        elif key == "M":
+        if key == "M":
             self.master.view.toggle_marked()
         elif key == "n":
             signals.status_prompt_onekey.send(
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 9b651dcc..457f3721 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -186,7 +186,8 @@ def default_keymap(km):
     km.add("F", "set console_focus_follow=toggle", context="flowlist")
     km.add("g", "view.go 0", context="flowlist")
     km.add("G", "view.go -1", context="flowlist")
-    km.add("l", "console.command 'cut.clip '", context="flowlist")
+    km.add("l", "console.command cut.clip ", context="flowlist")
+    km.add("L", "console.command view.load ", context="flowlist")
     km.add("m", "flow.mark.toggle @focus", context="flowlist")
     km.add("r", "replay.client @focus", context="flowlist")
     km.add("S", "console.command 'replay.server '")
diff --git a/mitmproxy/types/multidict.py b/mitmproxy/types/multidict.py
index c4f42580..bd9766a3 100644
--- a/mitmproxy/types/multidict.py
+++ b/mitmproxy/types/multidict.py
@@ -155,22 +155,6 @@ class _MultiDict(MutableMapping, metaclass=ABCMeta):
         else:
             return super().items()
 
-    def collect(self):
-        """
-            Returns a list of (key, value) tuples, where values are either
-            singular if there is only one matching item for a key, or a list
-            if there are more than one. The order of the keys matches the order
-            in the underlying fields list.
-        """
-        coll = []
-        for key in self:
-            values = self.get_all(key)
-            if len(values) == 1:
-                coll.append([key, values[0]])
-            else:
-                coll.append([key, values])
-        return coll
-
 
 class MultiDict(_MultiDict, serializable.Serializable):
     def __init__(self, fields=()):
diff --git a/setup.cfg b/setup.cfg
index 14502e20..993cad31 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -83,7 +83,6 @@ exclude =
     mitmproxy/proxy/root_context.py
     mitmproxy/proxy/server.py
     mitmproxy/stateobject.py
-    mitmproxy/types/multidict.py
     mitmproxy/utils/bits.py
     pathod/language/actions.py
     pathod/language/base.py
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 979f0aa1..eddcb04c 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -5,6 +5,7 @@ from mitmproxy.test import tflow
 from mitmproxy.addons import view
 from mitmproxy import flowfilter
 from mitmproxy import exceptions
+from mitmproxy import io
 from mitmproxy.test import taddons
 
 
@@ -130,10 +131,28 @@ def test_filter():
     assert len(v) == 4
 
 
-def test_load():
+def tdump(path, flows):
+    w = io.FlowWriter(open(path, "wb"))
+    for i in flows:
+        w.add(i)
+
+
+def test_load(tmpdir):
+    path = str(tmpdir.join("path"))
     v = view.View()
     with taddons.context() as tctx:
         tctx.master.addons.add(v)
+        tdump(
+            path,
+            [
+                tflow.tflow(resp=True),
+                tflow.tflow(resp=True)
+            ]
+        )
+        v.load_file(path)
+        assert len(v) == 2
+        v.load_file(path)
+        assert len(v) == 4
 
 
 def test_resolve():