Sketch out addons.View

The first iteration of a replacement for addons.State

3 files changed, 217 insertions, 0 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
new file mode 100644
index 00000000..f79dd64d
--- /dev/null
+++ b/mitmproxy/addons/view.py
@@ -0,0 +1,124 @@
+"""
+The View:
+
+- Keeps track of a store of flows
+- Maintains a filtered, ordered view onto that list of flows
+- Exposes various operations on flows in the store - notably intercept and
+  resume
+- Exposes a number of signals so the view can be monitored
+"""
+import collections
+import typing
+import datetime
+
+import blinker
+import sortedcontainers
+
+from mitmproxy import flow
+from mitmproxy import flowfilter
+
+
+def key_request_start(f: flow.Flow) -> datetime.datetime:
+    return f.request.timestamp_start or 0
+
+
+def key_request_method(f: flow.Flow) -> str:
+    return f.request.method
+
+
+matchall = flowfilter.parse(".")
+
+
+class View(collections.Sequence):
+    def __init__(self):
+        super().__init__()
+        self._store = {}
+        self.filter = matchall
+        self.order_key = key_request_start
+        self.order_reverse = False
+        self._view = sortedcontainers.SortedListWithKey(
+            key = self.order_key
+        )
+
+        # These signals broadcast events that affect the view. That is, an
+        # update to a flow in the store but not in the view does not trigger a
+        # signal.
+        self.sig_update = blinker.Signal()
+        self.sig_add = blinker.Signal()
+        self.sig_remove = blinker.Signal()
+        # Signals that the view should be refreshed completely
+        self.sig_refresh = blinker.Signal()
+
+    def __len__(self):
+        return len(self._view)
+
+    def __getitem__(self, offset) -> flow.Flow:
+        if self.order_reverse:
+            offset = -offset - 1
+        return self._view[offset]
+
+    def set_order(self, order_key: typing.Callable):
+        """
+            Sets the current view order.
+        """
+        self.order_key = order_key
+        newview = sortedcontainers.SortedListWithKey(key=order_key)
+        newview.update(self._view)
+        self._view = newview
+
+    def set_filter(self, flt: typing.Optional[flowfilter.TFilter]):
+        """
+            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)
+
+    def clear(self):
+        """
+            Clears both the state and view.
+        """
+        self._state.clear()
+        self._view.clear()
+
+    def add(self, f: flow.Flow):
+        """
+            Adds a flow to the state. If the flow already exists, it is
+            ignored.
+        """
+        if f.id not in self._store:
+            self._store[f.id] = f
+            if self.filter(f):
+                self._view.add(f)
+
+    def update(self, f: flow.Flow):
+        """
+            Updates a flow. If the flow is not in the state, it's ignored.
+        """
+        if f.id in self._store:
+            if self.filter(f):
+                if f not in self._view:
+                    self._view.add(f)
+            else:
+                try:
+                    self._view.remove(f)
+                except ValueError:
+                    pass
+
+    # Event handlers
+    def request(self, f):
+        self.add(f)
+
+    def intercept(self, f):
+        self.update(f)
+
+    def resume(self, f):
+        self.update(f)
+
+    def error(self, f):
+        self.update(f)
+
+    def response(self, f):
+        self.update(f)
diff --git a/setup.py b/setup.py
index 70ff8b5d..fd291973 100644
--- a/setup.py
+++ b/setup.py
@@ -83,6 +83,7 @@ setup(
         "urwid>=1.3.1, <1.4",
         "watchdog>=0.8.3, <0.9",
         "brotlipy>=0.5.1, <0.7",
+        "sortedcontainers>=1.5.4, <1.6",
     ],
     extras_require={
         ':sys_platform == "win32"': [
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
new file mode 100644
index 00000000..3e01d71f
--- /dev/null
+++ b/test/mitmproxy/addons/test_view.py
@@ -0,0 +1,92 @@
+from mitmproxy.addons import view
+from mitmproxy import flowfilter
+from .. import tutils
+
+
+def test_simple():
+    v = view.View()
+    f = tutils.tflow()
+    f.request.timestamp_start = 1
+    v.request(f)
+    assert list(v) == [f]
+    v.request(f)
+    assert list(v) == [f]
+    assert len(v._store) == 1
+
+    f2 = tutils.tflow()
+    f2.request.timestamp_start = 3
+    v.request(f2)
+    assert list(v) == [f, f2]
+    v.request(f2)
+    assert list(v) == [f, f2]
+    assert len(v._store) == 2
+
+    f3 = tutils.tflow()
+    f3.request.timestamp_start = 2
+    v.request(f3)
+    assert list(v) == [f, f3, f2]
+    v.request(f3)
+    assert list(v) == [f, f3, f2]
+    assert len(v._store) == 3
+
+
+def tft(*, method="get", start=0):
+    f = tutils.tflow()
+    f.request.method = method
+    f.request.timestamp_start = start
+    return f
+
+
+def test_filter():
+    v = view.View()
+    f = flowfilter.parse("~m get")
+    v.request(tft(method="get"))
+    v.request(tft(method="put"))
+    v.request(tft(method="get"))
+    v.request(tft(method="put"))
+    assert(len(v)) == 4
+    v.set_filter(f)
+    assert [i.request.method for i in v] == ["GET", "GET"]
+    assert len(v._store) == 4
+    v.set_filter(None)
+
+
+def test_order():
+    v = view.View()
+    v.request(tft(method="get", start=1))
+    v.request(tft(method="put", start=2))
+    v.request(tft(method="get", start=3))
+    v.request(tft(method="put", start=4))
+    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
+
+    v.set_order(view.key_request_method)
+    assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
+    v.order_reverse = True
+    assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"]
+
+    v.set_order(view.key_request_start)
+    assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1]
+
+    v.order_reverse = False
+    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
+
+
+def test_update():
+    v = view.View()
+    flt = flowfilter.parse("~m get")
+    v.set_filter(flt)
+
+    f = tft(method="get")
+    v.request(f)
+    assert f in v
+
+    f.request.method = "put"
+    v.update(f)
+    assert f not in v
+
+    f.request.method = "get"
+    v.update(f)
+    assert f in v
+
+    v.update(f)
+    assert f in v