From 9abfb1aac27e28619c58d40ec050cc46fda2b30f Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 30 Oct 2016 19:47:32 +1300 Subject: console: re-add follow focus --- mitmproxy/addons/view.py | 13 ++++++++++++- mitmproxy/tools/cmdline.py | 4 ++-- mitmproxy/tools/console/flowlist.py | 33 +++++++++++++++++---------------- mitmproxy/tools/console/master.py | 7 ++----- mitmproxy/tools/console/statusbar.py | 2 +- mitmproxy/tools/main.py | 2 +- test/mitmproxy/addons/test_view.py | 5 +++++ 7 files changed, 40 insertions(+), 26 deletions(-) diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index b4ba2315..9d38d94c 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -61,6 +61,8 @@ class View(collections.Sequence): self.show_marked = False self.order_key = key_request_start 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,6 +117,9 @@ 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 _refilter(self): self._view.clear() for i in self._store.values(): @@ -166,6 +171,8 @@ class View(collections.Sequence): self._store[f.id] = f if self.filter(f): self._view.add(f) + if self.focus_follow: + self.focus.flow = f self.sig_add.send(self, flow=f) def remove(self, f: mitmproxy.flow.Flow): @@ -186,6 +193,8 @@ class View(collections.Sequence): if self.filter(f): if f not in self._view: self._view.add(f) + if self.focus_follow: + self.focus.flow = f self.sig_add.send(self, flow=f) else: self.sig_update.send(self, flow=f) @@ -222,6 +231,8 @@ class View(collections.Sequence): ) 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) @@ -258,7 +269,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 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..77a8da30 100644 --- a/test/mitmproxy/addons/test_view.py +++ b/test/mitmproxy/addons/test_view.py @@ -7,6 +7,11 @@ from mitmproxy.test import taddons from .. import tutils +def test_keys(): + t = tflow.tflow(resp=True) + view.key_size(t) + + def test_simple(): v = view.View() f = tflow.tflow() -- cgit v1.2.3 From daf355bb4c7e7e1574d2977a49503eb6db2faa4a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 30 Oct 2016 21:09:08 +1300 Subject: console: add caching sort keys This is a tad complicated. 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. --- mitmproxy/addons/view.py | 137 +++++++++++++++++++++++++++---------- mitmproxy/master.py | 8 ++- 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: -- cgit v1.2.3