diff options
| -rw-r--r-- | mitmproxy/addons/view.py | 150 | ||||
| -rw-r--r-- | mitmproxy/master.py | 8 | ||||
| -rw-r--r-- | mitmproxy/tools/cmdline.py | 4 | ||||
| -rw-r--r-- | mitmproxy/tools/console/flowlist.py | 33 | ||||
| -rw-r--r-- | mitmproxy/tools/console/master.py | 7 | ||||
| -rw-r--r-- | mitmproxy/tools/console/statusbar.py | 2 | ||||
| -rw-r--r-- | mitmproxy/tools/main.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_view.py | 58 |
8 files changed, 171 insertions, 93 deletions
diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index b4ba2315..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,8 +103,18 @@ 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 + self._view = sortedcontainers.SortedListWithKey(key = self.order_key) # These signals broadcast events that affect the view. That is, an @@ -115,13 +169,23 @@ class View(collections.Sequence): def index(self, f: mitmproxy.flow.Flow) -> int: return self._rev(self._view.index(f)) + 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 @@ -155,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. @@ -165,7 +230,9 @@ 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) def remove(self, f: mitmproxy.flow.Flow): @@ -173,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): """ @@ -185,9 +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: @@ -210,32 +284,31 @@ 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: + self.focus_follow = opts.focus_follow 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) @@ -258,7 +331,7 @@ class Focus: return self._flow @flow.setter - def flow(self, f: mitmproxy.flow.Flow): + def flow(self, f: typing.Optional[mitmproxy.flow.Flow]): if f is not None and f not in self.view: raise ValueError("Attempt to set focus to flow not in view") self._flow = f @@ -305,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) @@ -321,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/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index cf6e1d35..e4b29d0f 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -784,8 +784,8 @@ def mitmproxy(): ) parser.add_argument( "--follow", - action="store_true", dest="follow", - help="Follow flow list." + action="store_true", dest="focus_follow", + help="Focus follows new flows." ) parser.add_argument( "--order", diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index a0a3dc94..76545893 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -175,18 +175,6 @@ class FlowItem(urwid.WidgetWrap): signals.flowlist_change.send(self) elif key == "M": self.master.view.toggle_marked() - elif key == "o": - orders = [(i[1], i[0]) for i in view.orders] - lookup = dict([(i[0], i[1]) for i in view.orders]) - - def change_order(k): - self.master.options.order = lookup[k] - - signals.status_prompt_onekey.send( - prompt = "Order", - keys = orders, - callback = change_order - ) elif key == "r": try: self.master.replay_request(self.flow) @@ -220,9 +208,6 @@ class FlowItem(urwid.WidgetWrap): for f in self.master.view: f.marked = False signals.flowlist_change.send(self) - elif key == "v": - val = not self.master.options.order_reversed - self.master.options.order_reversed = val elif key == "V": if not self.flow.modified(): signals.status_message.send(message="Flow not modified.") @@ -389,8 +374,24 @@ class FlowListBox(urwid.ListBox): keys = common.METHOD_OPTIONS, callback = self.get_method ) + elif key == "o": + orders = [(i[1], i[0]) for i in view.orders] + lookup = dict([(i[0], i[1]) for i in view.orders]) + + def change_order(k): + self.master.options.order = lookup[k] + + signals.status_prompt_onekey.send( + prompt = "Order", + keys = orders, + callback = change_order + ) elif key == "F": - self.master.toggle_follow_flows() + o = self.master.options + o.focus_follow = not o.focus_follow + elif key == "v": + val = not self.master.options.order_reversed + self.master.options.order_reversed = val elif key == "W": if self.master.options.outfile: self.master.options.outfile = None diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 932dc151..4c5e8c8c 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -44,25 +44,23 @@ class Options(mitmproxy.options.Options): self, *, # all args are keyword-only. eventlog: bool = False, - follow: bool = False, + focus_follow: bool = False, intercept: Optional[str] = None, filter: Optional[str] = None, palette: Optional[str] = None, palette_transparent: bool = False, no_mouse: bool = False, - follow_focus: bool = False, order: Optional[str] = None, order_reversed: bool = False, **kwargs ): self.eventlog = eventlog - self.follow = follow + self.focus_follow = focus_follow self.intercept = intercept self.filter = filter self.palette = palette self.palette_transparent = palette_transparent self.no_mouse = no_mouse - self.follow_focus = follow_focus self.order = order self.order_reversed = order_reversed super().__init__(**kwargs) @@ -83,7 +81,6 @@ class ConsoleMaster(master.Master): self.palette_transparent = options.palette_transparent self.logbuffer = urwid.SimpleListWalker([]) - self.follow = options.follow self.view_stack = [] diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index b358f711..e292cbd7 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -199,7 +199,7 @@ class StatusBar(urwid.WidgetWrap): opts.append("killextra") if self.master.options.no_upstream_cert: opts.append("no-upstream-cert") - if self.master.options.follow_focus: + if self.master.options.focus_follow: opts.append("following") if self.master.options.stream_large_bodies: opts.append( diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 772841f5..08e5f2ca 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -69,7 +69,7 @@ def mitmproxy(args=None): # pragma: no cover console_options.palette = args.palette console_options.palette_transparent = args.palette_transparent console_options.eventlog = args.eventlog - console_options.follow = args.follow + console_options.focus_follow = args.focus_follow console_options.intercept = args.intercept console_options.filter = args.filter console_options.no_mouse = args.no_mouse diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py index 750e8469..63df8307 100644 --- a/test/mitmproxy/addons/test_view.py +++ b/test/mitmproxy/addons/test_view.py @@ -7,6 +7,20 @@ from mitmproxy.test import taddons from .. import tutils +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(): v = view.View() f = tflow.tflow() @@ -69,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] - v.set_order(view.key_request_start) - assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1] + 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"] + + 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(): @@ -256,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 @@ -265,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: |