diff options
| -rw-r--r-- | mitmproxy/addons/view.py | 137 | ||||
| -rw-r--r-- | mitmproxy/master.py | 8 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_view.py | 59 |
3 files changed, 134 insertions, 70 deletions
diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index 9d38d94c..c09ff454 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -19,39 +19,83 @@ import mitmproxy.flow from mitmproxy import flowfilter from mitmproxy import exceptions +# The underlying sorted list implementation expects the sort key to be stable +# for the lifetime of the object. However, if we sort by size, for instance, +# the sort order changes as the flow progresses through its lifecycle. We +# address this through two means: +# +# - Let order keys cache the sort value by flow ID. +# +# - Add a facility to refresh items in the list by removing and re-adding them +# when they are updated. + + +class _OrderKey: + def __init__(self, view): + self.view = view -def key_request_start(f: mitmproxy.flow.Flow) -> datetime.datetime: - return f.request.timestamp_start or 0 + def generate(self, f: mitmproxy.flow.Flow) -> typing.Any: + pass + def refresh(self, f): + k = self._key() + old = self.view.settings[f][k] + new = self.generate(f) + if old != new: + self.view._view.remove(f) + self.view.settings[f][k] = new + self.view._view.add(f) + self.view.sig_refresh.send(self.view) -def key_request_method(f: mitmproxy.flow.Flow) -> str: - return f.request.method + def _key(self): + return "_order_%s" % id(self) + def __call__(self, f): + k = self._key() + s = self.view.settings[f] + if k in s: + return s[k] + val = self.generate(f) + s[k] = val + return val -def key_request_url(f: mitmproxy.flow.Flow) -> str: - return f.request.url +class OrderRequestStart(_OrderKey): + def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + return f.request.timestamp_start or 0 -def key_size(f: mitmproxy.flow.Flow) -> int: - s = 0 - if f.request.raw_content: - s += len(f.request.raw_content) - if f.response and f.response.raw_content: - s += len(f.response.raw_content) - return s +class OrderRequestMethod(_OrderKey): + def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + return f.request.method -orders = [ - ("t", "time", key_request_start), - ("m", "method", key_request_method), - ("u", "url", key_request_url), - ("z", "size", key_size), -] + +class OrderRequestURL(_OrderKey): + def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + return f.request.url + + +class OrderKeySize(_OrderKey): + def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime: + s = 0 + if f.request.raw_content: + s += len(f.request.raw_content) + if f.response and f.response.raw_content: + s += len(f.response.raw_content) + return s matchall = flowfilter.parse(".") +orders = [ + ("t", "time"), + ("m", "method"), + ("u", "url"), + ("z", "size"), +] + + class View(collections.Sequence): def __init__(self): super().__init__() @@ -59,7 +103,15 @@ class View(collections.Sequence): self.filter = matchall # Should we show only marked flows? self.show_marked = False - self.order_key = key_request_start + + self.default_order = OrderRequestStart(self) + self.orders = dict( + time = self.default_order, + method = OrderRequestMethod(self), + url = OrderRequestURL(self), + size = OrderKeySize(self), + ) + self.order_key = self.default_order self.order_reversed = False self.focus_follow = False @@ -120,13 +172,20 @@ class View(collections.Sequence): def __contains__(self, f: mitmproxy.flow.Flow) -> bool: return self._view.__contains__(f) + def _order_key_name(self): + return "_order_%s" % id(self.order_key) + + def _base_add(self, f): + self.settings[f][self._order_key_name()] = self.order_key(f) + self._view.add(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._base_add(i) self.sig_refresh.send(self) # API @@ -160,9 +219,10 @@ class View(collections.Sequence): """ self._store.clear() self._view.clear() + self.settings.clear() self.sig_refresh.send(self) - def add(self, f: mitmproxy.flow.Flow): + def add(self, f: mitmproxy.flow.Flow) -> bool: """ Adds a flow to the state. If the flow already exists, it is ignored. @@ -170,7 +230,7 @@ class View(collections.Sequence): if f.id not in self._store: self._store[f.id] = f if self.filter(f): - self._view.add(f) + self._base_add(f) if self.focus_follow: self.focus.flow = f self.sig_add.send(self, flow=f) @@ -180,10 +240,10 @@ class View(collections.Sequence): Removes the flow from the underlying store and the view. """ if f.id in self._store: - del self._store[f.id] if f in self._view: self._view.remove(f) self.sig_remove.send(self, flow=f) + del self._store[f.id] def update(self, f: mitmproxy.flow.Flow): """ @@ -192,11 +252,16 @@ class View(collections.Sequence): if f.id in self._store: if self.filter(f): if f not in self._view: - self._view.add(f) + self._base_add(f) if self.focus_follow: self.focus.flow = f self.sig_add.send(self, flow=f) else: + # This is a tad complicated. The sortedcontainers + # implementation assumes that the order key is stable. If + # it changes mid-way Very Bad Things happen. We detect when + # this happens, and re-fresh the item. + self.order_key.refresh(f) self.sig_update.send(self, flow=f) else: try: @@ -219,16 +284,13 @@ class View(collections.Sequence): self.set_filter(filt) if "order" in updated: if opts.order is None: - self.set_order(key_request_start) + self.set_order(self.default_order) else: - for _, name, func in orders: - if name == opts.order: - self.set_order(func) - break - else: + if opts.order not in self.orders: raise exceptions.OptionsError( "Unknown flow order: %s" % opts.order ) + self.set_order(self.orders[opts.order]) if "order_reversed" in updated: self.set_reversed(opts.order_reversed) if "focus_follow" in updated: @@ -237,16 +299,16 @@ class View(collections.Sequence): def request(self, f): self.add(f) - def intercept(self, f): + def error(self, f): self.update(f) - def resume(self, f): + def response(self, f): self.update(f) - def error(self, f): + def intercept(self, f): self.update(f) - def response(self, f): + def resume(self, f): self.update(f) @@ -316,6 +378,9 @@ class Settings(collections.Mapping): view.sig_remove.connect(self._sig_remove) view.sig_refresh.connect(self._sig_refresh) + def clear(self): + self.values.clear() + def __iter__(self) -> typing.Iterable: return iter(self.values) @@ -332,6 +397,6 @@ class Settings(collections.Mapping): del self.values[flow.id] def _sig_refresh(self, view): - for fid in self.values.keys(): + for fid in list(self.values.keys()): if fid not in view._store: del self.values[fid] diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 31ce17a3..ffbfb0cb 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -92,10 +92,14 @@ class Master: try: mtype, obj = self.event_queue.get(timeout=timeout) if mtype not in events.Events: - raise exceptions.ControlException("Unknown event %s" % repr(mtype)) + raise exceptions.ControlException( + "Unknown event %s" % repr(mtype) + ) handle_func = getattr(self, mtype) if not callable(handle_func): - raise exceptions.ControlException("Handler %s not callable" % mtype) + raise exceptions.ControlException( + "Handler %s not callable" % mtype + ) if not handle_func.__dict__.get("__handler"): raise exceptions.ControlException( "Handler function %s is not decorated with controller.handler" % ( diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py index 77a8da30..63df8307 100644 --- a/test/mitmproxy/addons/test_view.py +++ b/test/mitmproxy/addons/test_view.py @@ -7,9 +7,18 @@ from mitmproxy.test import taddons from .. import tutils -def test_keys(): - t = tflow.tflow(resp=True) - view.key_size(t) +class Options(options.Options): + def __init__( + self, *, + filter=None, + order=None, + order_reversed=False, + **kwargs + ): + self.filter = filter + self.order = order + self.order_reversed = order_reversed + super().__init__(**kwargs) def test_simple(): @@ -74,22 +83,23 @@ def test_filter(): 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.set_reversed(True) - assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"] + with taddons.context(options=Options()) as tctx: + 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] + + tctx.configure(v, order="method") + assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"] + v.set_reversed(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] + tctx.configure(v, order="time") + assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1] - v.set_reversed(False) - assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4] + v.set_reversed(False) + assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4] def test_reversed(): @@ -261,7 +271,6 @@ def test_settings(): 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 @@ -270,20 +279,6 @@ def test_settings(): assert not v.settings.keys() -class Options(options.Options): - def __init__( - self, *, - filter=None, - order=None, - order_reversed=False, - **kwargs - ): - self.filter = filter - self.order = order - self.order_reversed = order_reversed - super().__init__(**kwargs) - - def test_configure(): v = view.View() with taddons.context(options=Options()) as tctx: |