Merge pull request #1692 from cortesi/floworder

console: o for flow order

8 files changed, 88 insertions, 10 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index d62221c5..b4ba2315 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -28,6 +28,27 @@ def key_request_method(f: mitmproxy.flow.Flow) -> str:
     return f.request.method
 
 
+def key_request_url(f: mitmproxy.flow.Flow) -> str:
+    return f.request.url
+
+
+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
+
+
+orders = [
+    ("t", "time", key_request_start),
+    ("m", "method", key_request_method),
+    ("u", "url", key_request_url),
+    ("z", "size", key_size),
+]
+
+
 matchall = flowfilter.parse(".")
 
 
@@ -108,8 +129,8 @@ class View(collections.Sequence):
         self.show_marked = not self.show_marked
         self._refilter()
 
-    def toggle_reversed(self):
-        self.order_reversed = not self.order_reversed
+    def set_reversed(self, value: bool):
+        self.order_reversed = value
         self.sig_refresh.send(self)
 
     def set_order(self, order_key: typing.Callable):
@@ -178,15 +199,29 @@ class View(collections.Sequence):
 
     # Event handlers
     def configure(self, opts, updated):
-        filt = None
         if "filter" in updated:
+            filt = None
             if opts.filter:
                 filt = flowfilter.parse(opts.filter)
                 if not filt:
                     raise exceptions.OptionsError(
                         "Invalid interception filter: %s" % opts.filter
                     )
-        self.set_filter(filt)
+            self.set_filter(filt)
+        if "order" in updated:
+            if opts.order is None:
+                self.set_order(key_request_start)
+            else:
+                for _, name, func in orders:
+                    if name == opts.order:
+                        self.set_order(func)
+                        break
+                else:
+                    raise exceptions.OptionsError(
+                        "Unknown flow order: %s" % opts.order
+                    )
+        if "order_reversed" in updated:
+            self.set_reversed(opts.order_reversed)
 
     def request(self, f):
         self.add(f)
diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py
index 55adb7fa..cf6e1d35 100644
--- a/mitmproxy/tools/cmdline.py
+++ b/mitmproxy/tools/cmdline.py
@@ -8,6 +8,7 @@ from mitmproxy import platform
 from mitmproxy.utils import human
 from mitmproxy.net import tcp
 from mitmproxy import version
+from mitmproxy.addons import view
 
 
 class ParseException(Exception):
@@ -787,6 +788,12 @@ def mitmproxy():
         help="Follow flow list."
     )
     parser.add_argument(
+        "--order",
+        type=str, dest="order", default=None,
+        choices=[o[1] for o in view.orders],
+        help="Flow sort order."
+    )
+    parser.add_argument(
         "--no-mouse",
         action="store_true", dest="no_mouse",
         help="Disable mouse interaction."
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 677194b3..a0a3dc94 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -4,6 +4,7 @@ import mitmproxy.net.http.url
 from mitmproxy import exceptions
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
+from mitmproxy.addons import view
 from mitmproxy import export
 
 
@@ -24,9 +25,11 @@ def _mkhelp():
         ("m", "toggle flow mark"),
         ("M", "toggle marked flow view"),
         ("n", "create a new request"),
+        ("o", "set flow order"),
         ("r", "replay request"),
         ("S", "server replay request/s"),
         ("U", "unmark all marked flows"),
+        ("v", "reverse flow order"),
         ("V", "revert changes to request"),
         ("w", "save flows "),
         ("W", "stream flows to file"),
@@ -172,6 +175,18 @@ 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)
@@ -205,6 +220,9 @@ 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.")
diff --git a/mitmproxy/tools/console/help.py b/mitmproxy/tools/console/help.py
index dda8bfbc..282f374d 100644
--- a/mitmproxy/tools/console/help.py
+++ b/mitmproxy/tools/console/help.py
@@ -48,7 +48,7 @@ class HelpView(urwid.ListBox):
         text.append(urwid.Text([("head", "\n\nGlobal keys:\n")]))
         keys = [
             ("i", "set interception pattern"),
-            ("o", "options"),
+            ("O", "options"),
             ("q", "quit / return to previous page"),
             ("Q", "quit without confirm prompt"),
             ("R", "replay of requests/responses from file"),
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 6d94611e..932dc151 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -51,6 +51,8 @@ class Options(mitmproxy.options.Options):
             palette_transparent: bool = False,
             no_mouse: bool = False,
             follow_focus: bool = False,
+            order: Optional[str] = None,
+            order_reversed: bool = False,
             **kwargs
     ):
         self.eventlog = eventlog
@@ -61,6 +63,8 @@ class Options(mitmproxy.options.Options):
         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)
 
 
diff --git a/mitmproxy/tools/console/window.py b/mitmproxy/tools/console/window.py
index a01c59bc..1c962e2f 100644
--- a/mitmproxy/tools/console/window.py
+++ b/mitmproxy/tools/console/window.py
@@ -91,7 +91,7 @@ class Window(urwid.Frame):
                 text = self.master.options.intercept,
                 callback = self.master.options.setter("intercept")
             )
-        elif k == "o":
+        elif k == "O":
             self.master.view_options()
         elif k == "Q":
             raise urwid.ExitMainLoop
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index f88bee6c..772841f5 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -73,6 +73,7 @@ def mitmproxy(args=None):  # pragma: no cover
         console_options.intercept = args.intercept
         console_options.filter = args.filter
         console_options.no_mouse = args.no_mouse
+        console_options.order = args.order
 
         server = process_options(parser, console_options, args)
         m = console.master.ConsoleMaster(console_options, server)
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index c7a9821a..750e8469 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -77,13 +77,13 @@ def test_order():
 
     v.set_order(view.key_request_method)
     assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
-    v.toggle_reversed()
+    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]
 
-    v.toggle_reversed()
+    v.set_reversed(False)
     assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
 
 
@@ -92,7 +92,7 @@ def test_reversed():
     v.request(tft(start=1))
     v.request(tft(start=2))
     v.request(tft(start=3))
-    v.toggle_reversed()
+    v.set_reversed(True)
 
     assert v[0].request.timestamp_start == 3
     assert v[-1].request.timestamp_start == 1
@@ -266,8 +266,16 @@ def test_settings():
 
 
 class Options(options.Options):
-    def __init__(self, *, filter=None, **kwargs):
+    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)
 
 
@@ -276,3 +284,8 @@ def test_configure():
     with taddons.context(options=Options()) as tctx:
         tctx.configure(v, filter="~q")
         tutils.raises("invalid interception filter", tctx.configure, v, filter="~~")
+
+        tctx.configure(v, order="method")
+        tutils.raises("unknown flow order", tctx.configure, v, order="no")
+
+        tctx.configure(v, order_reversed=True)