implement a toggle for viewing marked flows only in console

3 files changed, 44 insertions, 0 deletions

diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py
index 1dd032be..4541a6af 100644
--- a/mitmproxy/console/__init__.py
+++ b/mitmproxy/console/__init__.py
@@ -33,6 +33,8 @@ class ConsoleState(flow.State):
         self.default_body_view = contentviews.get("Auto")
         self.flowsettings = weakref.WeakKeyDictionary()
         self.last_search = None
+        self.last_filter = None
+        self.mark_filter = False
 
     def __setattr__(self, name, value):
         self.__dict__[name] = value
@@ -106,6 +108,37 @@ class ConsoleState(flow.State):
         self.set_focus(self.focus)
         return ret
 
+    def filter_marked(self, m):
+        def actual_func(x):
+            if x.id in m:
+                return True
+            return False
+        return actual_func
+
+    def enable_marked_filter(self):
+        self.last_filter = self.limit_txt
+        marked_flows = []
+        for f in self.flows:
+            if self.flow_marked(f):
+                marked_flows.append(f.id)
+        if len(marked_flows) > 0:
+            f = self.filter_marked(marked_flows)
+            self.view._close()
+            self.view = flow.FlowView(self.flows, f)
+            self.focus = 0
+            self.set_focus(self.focus)
+            self.mark_filter = True
+
+    def disable_marked_filter(self):
+        if self.last_filter is None:
+            self.view = flow.FlowView(self.flows, None)
+        else:
+            self.set_limit(self.last_filter)
+        self.focus = 0
+        self.set_focus(self.focus)
+        self.last_filter = None
+        self.mark_filter = False
+
     def clear(self):
         marked_flows = []
         for f in self.flows:
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 78b30231..5672c0d9 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -21,6 +21,7 @@ def _mkhelp():
         ("l", "set limit filter pattern"),
         ("L", "load saved flows"),
         ("m", "toggle flow mark"),
+        ("M", "toggle marked flow view"),
         ("n", "create a new request"),
         ("P", "copy flow to clipboard"),
         ("r", "replay request"),
@@ -197,6 +198,12 @@ class ConnectionItem(urwid.WidgetWrap):
             else:
                 self.state.set_flow_marked(self.flow, True)
             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)
         elif key == "r":
             r = self.master.replay_request(self.flow)
             if r:
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index 4cc63a54..96840a20 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -163,6 +163,10 @@ class StatusBar(urwid.WidgetWrap):
             r.append("[")
             r.append(("heading_key", "l"))
             r.append(":%s]" % self.master.state.limit_txt)
+        if self.master.state.mark_filter:
+            r.append("[")
+            r.append(("heading_key", "Marked Flows"))
+            r.append("]")
         if self.master.stickycookie_txt:
             r.append("[")
             r.append(("heading_key", "t"))