addons.view: Add first-class support for marking

Marking now works differently - it's no longer just another filter, it's
applied in addition to the filter. This means you can apply a filter, mark some
flows, and then toggle between marked and unmarked flows matching the filter.

I'm leaving the filter for marked flows intact - it will come in handy in other
situations.

4 files changed, 28 insertions, 103 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 9cde0345..3f98667a 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -36,6 +36,8 @@ class View(collections.Sequence):
         super().__init__()
         self._store = {}
         self.filter = matchall
+        # Should we show only marked flows?
+        self.show_marked = False
         self.order_key = key_request_start
         self.order_reversed = False
         self._view = sortedcontainers.SortedListWithKey(key = self.order_key)
@@ -92,7 +94,19 @@ class View(collections.Sequence):
     def index(self, f: mitmproxy.flow.Flow) -> int:
         return self._rev(self._view.index(f))
 
+    def _refilter(self):
+        self._view.clear()
+        for i in self._store.values():
+            if self.show_marked and not i.marked:
+                continue
+            if self.filter(i):
+                self._view.add(i)
+        self.sig_refresh.send(self)
+
     # API
+    def toggle_marked(self):
+        self.show_marked = not self.show_marked
+        self._refilter()
 
     def toggle_reversed(self):
         self.order_reversed = not self.order_reversed
@@ -112,11 +126,7 @@ class View(collections.Sequence):
             Sets the current view filter.
         """
         self.filter = flt or matchall
-        self._view.clear()
-        for i in self._store.values():
-            if self.filter(i):
-                self._view.add(i)
-        self.sig_refresh.send(self)
+        self._refilter()
 
     def clear(self):
         """
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 32fe0f7d..1d39c172 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -174,11 +174,7 @@ class ConnectionItem(urwid.WidgetWrap):
             self.flow.marked = not self.flow.marked
             signals.flowlist_change.send(self)
         elif key == "M":
-            if self.state.mark_filter:
-                self.state.disable_marked_filter()
-            else:
-                self.state.enable_marked_filter()
-            signals.flowlist_change.send(self)
+            self.master.view.toggle_marked()
         elif key == "r":
             try:
                 self.master.replay_request(self.flow)
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index e8d6fb07..c7a9821a 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -58,6 +58,14 @@ def test_filter():
     assert len(v._store) == 4
     v.set_filter(None)
 
+    assert len(v) == 4
+    v[1].marked = True
+    v.toggle_marked()
+    assert len(v) == 1
+    assert v[0].marked
+    v.toggle_marked()
+    assert len(v) == 4
+
 
 def test_order():
     v = view.View()
diff --git a/test/mitmproxy/console/test_master.py b/test/mitmproxy/console/test_master.py
index fe552062..eb840187 100644
--- a/test/mitmproxy/console/test_master.py
+++ b/test/mitmproxy/console/test_master.py
@@ -1,5 +1,3 @@
-import gc
-
 from mitmproxy.test import tflow
 import mitmproxy.test.tutils
 from mitmproxy.tools import console
@@ -9,93 +7,6 @@ from mitmproxy.tools.console import common
 from .. import mastertest
 
 
-class TestConsoleState:
-
-    def test_flow(self):
-        """
-            normal flow:
-
-                connect -> request -> response
-        """
-        c = console.master.ConsoleState()
-        f = self._add_request(c)
-        assert f in c.flows
-        assert c.get_focus() == (f, 0)
-
-    def test_focus(self):
-        """
-            normal flow:
-
-                connect -> request -> response
-        """
-        c = console.master.ConsoleState()
-        f = self._add_request(c)
-
-        assert c.get_focus() == (f, 0)
-        assert c.get_from_pos(0) == (f, 0)
-        assert c.get_from_pos(1) == (None, None)
-        assert c.get_next(0) == (None, None)
-
-        f2 = self._add_request(c)
-        assert c.get_focus() == (f, 0)
-        assert c.get_next(0) == (f2, 1)
-        assert c.get_prev(1) == (f, 0)
-        assert c.get_next(1) == (None, None)
-
-        c.set_focus(0)
-        assert c.get_focus() == (f, 0)
-        c.set_focus(-1)
-        assert c.get_focus() == (f, 0)
-        c.set_focus(2)
-        assert c.get_focus() == (f2, 1)
-
-        c.delete_flow(f2)
-        assert c.get_focus() == (f, 0)
-        c.delete_flow(f)
-        assert c.get_focus() == (None, None)
-
-    def _add_request(self, state):
-        f = tflow.tflow()
-        return state.add_flow(f)
-
-    def _add_response(self, state):
-        f = self._add_request(state)
-        f.response = mitmproxy.test.tutils.tresp()
-        state.update_flow(f)
-
-    def test_add_response(self):
-        c = console.master.ConsoleState()
-        f = self._add_request(c)
-        f.response = mitmproxy.test.tutils.tresp()
-        c.focus = None
-        c.update_flow(f)
-
-    def test_focus_view(self):
-        c = console.master.ConsoleState()
-        self._add_request(c)
-        self._add_response(c)
-        self._add_request(c)
-        self._add_response(c)
-        self._add_request(c)
-        self._add_response(c)
-        assert not c.set_view_filter("~s")
-        assert len(c.view) == 3
-        assert c.focus == 0
-
-    def test_settings(self):
-        c = console.master.ConsoleState()
-        f = self._add_request(c)
-        c.add_flow_setting(f, "foo", "bar")
-        assert c.get_flow_setting(f, "foo") == "bar"
-        assert c.get_flow_setting(f, "oink") is None
-        assert c.get_flow_setting(f, "oink", "foo") == "foo"
-        assert len(c.flowsettings) == 1
-        c.delete_flow(f)
-        del f
-        gc.collect()
-        assert len(c.flowsettings) == 0
-
-
 def test_format_keyvals():
     assert common.format_keyvals(
         [
@@ -123,17 +34,17 @@ class TestMaster(mastertest.MasterTest):
         m = self.mkmaster()
         for i in (1, 2, 3):
             self.dummy_cycle(m, 1, b"")
-            assert len(m.state.flows) == i
+            assert len(m.view) == i
 
     def test_intercept(self):
         """regression test for https://github.com/mitmproxy/mitmproxy/issues/1605"""
         m = self.mkmaster(intercept="~b bar")
         f = tflow.tflow(req=mitmproxy.test.tutils.treq(content=b"foo"))
         m.request(f)
-        assert not m.state.flows[0].intercepted
+        assert not m.view[0].intercepted
         f = tflow.tflow(req=mitmproxy.test.tutils.treq(content=b"bar"))
         m.request(f)
-        assert m.state.flows[1].intercepted
+        assert m.view[1].intercepted
         f = tflow.tflow(resp=mitmproxy.test.tutils.tresp(content=b"bar"))
         m.request(f)
-        assert m.state.flows[2].intercepted
+        assert m.view[2].intercepted