diff options
| -rw-r--r-- | mitmproxy/addons/view.py | 124 | ||||
| -rw-r--r-- | setup.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_view.py | 92 |
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) @@ -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 |