addons.view: flow settings

Add a flow settings mechanism, enable focus and settings unilaterally.

2 files changed, 51 insertions, 1 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index d8d6e853..e2224c58 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -4,7 +4,9 @@ The View:
 - Keeps track of a store of flows
 - Maintains a filtered, ordered view onto that list of flows
 - Exposes a number of signals so the view can be monitored
-- Has an associated class that tracks focus within the view
+- Tracks focus within the view
+- Exposes a settings store for flows that automatically expires if the flow is
+  removed from the store.
 """
 import collections
 import typing
@@ -46,6 +48,9 @@ class View(collections.Sequence):
         # Signals that the view should be refreshed completely
         self.sig_refresh = blinker.Signal()
 
+        self.focus = Focus(self)
+        self.settings = Settings(self)
+
     def _rev(self, idx: int) -> int:
         """
             Reverses an index, if needed
@@ -211,6 +216,8 @@ class Focus:
     def _sig_refresh(self, view):
         if len(view) == 0:
             self.focusflow = None
+        elif self.focusflow is None:
+            self.focusflow = view[0]
         elif self.focusflow not in view:
             self.focusflow = view[self._nearest(self.focusflow, view)]
 
@@ -218,3 +225,31 @@ class Focus:
         # We only have to act if we don't have a focus element
         if not self.focusflow:
             self.focusflow = flow
+
+
+class Settings(collections.Mapping):
+    def __init__(self, view: View) -> None:
+        self.view = view
+        self.values = {}
+        view.sig_remove.connect(self._sig_remove)
+        view.sig_refresh.connect(self._sig_refresh)
+
+    def __iter__(self) -> typing.Iterable:
+        return iter(self.values)
+
+    def __len__(self) -> int:
+        return len(self.values)
+
+    def __getitem__(self, f: flow.Flow) -> dict:
+        if f.id not in self.view._store:
+            raise KeyError
+        return self.values.setdefault(f.id, {})
+
+    def _sig_remove(self, view, flow):
+        if flow.id in self.values:
+            del self.values[flow.id]
+
+    def _sig_refresh(self, view):
+        for fid in self.values.keys():
+            if fid not in view._store:
+                del self.values[fid]
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 56372749..1404a78a 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -233,3 +233,18 @@ def test_focus():
     filt = flowfilter.parse("~m oink")
     v.set_filter(filt)
     assert f.index is None
+
+
+def test_settings():
+    v = view.View()
+    f = tft()
+
+    tutils.raises(KeyError, v.settings.__getitem__, f)
+    v.add(f)
+    assert v.settings[f] == {}
+    v.settings[f]["foo"] = "bar"
+    assert v.settings[f]["foo"] == "bar"
+    assert len(list(v.settings)) == 1
+    v.remove(f)
+    tutils.raises(KeyError, v.settings.__getitem__, f)
+    assert not v.settings.keys()