console: extract grideditor keybindings

Also formalise the LayoutWidget interface a bit

12 files changed, 191 insertions, 148 deletions

diff --git a/mitmproxy/tools/console/commands.py b/mitmproxy/tools/console/commands.py
index ca6e6dfe..5ab55e48 100644
--- a/mitmproxy/tools/console/commands.py
+++ b/mitmproxy/tools/console/commands.py
@@ -2,6 +2,7 @@ import urwid
 import blinker
 import textwrap
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 from mitmproxy.tools.console import signals
 
 HELP_HEIGHT = 5
@@ -151,7 +152,7 @@ class CommandHelp(urwid.Frame):
         self.set_body(self.widget(txt))
 
 
-class Commands(urwid.Pile):
+class Commands(urwid.Pile, layoutwidget.LayoutWidget):
     title = "Commands"
     keyctx = "commands"
 
diff --git a/mitmproxy/tools/console/eventlog.py b/mitmproxy/tools/console/eventlog.py
index c56b80d3..1e56a05a 100644
--- a/mitmproxy/tools/console/eventlog.py
+++ b/mitmproxy/tools/console/eventlog.py
@@ -1,5 +1,6 @@
 import urwid
 from mitmproxy.tools.console import signals
+from mitmproxy.tools.console import layoutwidget
 
 EVENTLOG_SIZE = 10000
 
@@ -8,7 +9,7 @@ class LogBufferWalker(urwid.SimpleListWalker):
     pass
 
 
-class EventLog(urwid.ListBox):
+class EventLog(urwid.ListBox, layoutwidget.LayoutWidget):
     keyctx = "eventlog"
     title = "Events"
 
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index cb0aa1bd..e2d9a571 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -1,6 +1,7 @@
 import urwid
 
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 import mitmproxy.tools.console.master # noqa
 
 
@@ -109,7 +110,7 @@ class FlowListWalker(urwid.ListWalker):
         return f, pos
 
 
-class FlowListBox(urwid.ListBox):
+class FlowListBox(urwid.ListBox, layoutwidget.LayoutWidget):
     title = "Flows"
     keyctx = "flowlist"
 
diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py
index 58308efb..568c5e43 100644
--- a/mitmproxy/tools/console/flowview.py
+++ b/mitmproxy/tools/console/flowview.py
@@ -8,6 +8,7 @@ import urwid
 from mitmproxy import contentviews
 from mitmproxy import http
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 from mitmproxy.tools.console import flowdetailview
 from mitmproxy.tools.console import searchable
 from mitmproxy.tools.console import signals
@@ -274,7 +275,7 @@ class FlowDetails(tabs.Tabs):
         return self._w.keypress(size, key)
 
 
-class FlowView(urwid.Frame):
+class FlowView(urwid.Frame, layoutwidget.LayoutWidget):
     keyctx = "flowview"
     title = "Flow Details"
 
diff --git a/mitmproxy/tools/console/grideditor/base.py b/mitmproxy/tools/console/grideditor/base.py
index 34af7e7b..722fce55 100644
--- a/mitmproxy/tools/console/grideditor/base.py
+++ b/mitmproxy/tools/console/grideditor/base.py
@@ -12,6 +12,7 @@ from typing import Set # noqa
 import urwid
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
+from mitmproxy.tools.console import layoutwidget
 import mitmproxy.tools.console.master # noqa
 
 FOOTER = [
@@ -58,6 +59,7 @@ class Column(metaclass=abc.ABCMeta):
 
 
 class GridRow(urwid.WidgetWrap):
+
     def __init__(
             self,
             focused: Optional[int],
@@ -253,6 +255,7 @@ FIRST_WIDTH_MIN = 20
 
 
 class BaseGridEditor(urwid.WidgetWrap):
+    keyctx = "grideditor"
 
     def __init__(
             self,
@@ -346,14 +349,6 @@ class BaseGridEditor(urwid.WidgetWrap):
             self.walker.left()
         elif key == "right":
             self.walker.right()
-        elif key == "tab":
-            self.walker.tab_next()
-        elif key == "a":
-            self.walker.add()
-        elif key == "A":
-            self.walker.insert()
-        elif key == "d":
-            self.walker.delete_focus()
         elif column.keypress(key, self) and not self.handle_key(key):
             return self._w.keypress(size, key)
 
@@ -408,10 +403,23 @@ class BaseGridEditor(urwid.WidgetWrap):
         )
         return text
 
+    def cmd_next(self):
+        self.walker.tab_next()
+
+    def cmd_add(self):
+        self.walker.add()
+
+    def cmd_insert(self):
+        self.walker.insert()
+
+    def cmd_delete(self):
+        self.walker.delete_focus()
+
 
-class GridEditor(urwid.WidgetWrap):
+class GridEditor(BaseGridEditor):
     title = None  # type: str
     columns = None  # type: Sequence[Column]
+    keyctx = "grideditor"
 
     def __init__(
             self,
@@ -423,16 +431,16 @@ class GridEditor(urwid.WidgetWrap):
     ) -> None:
         super().__init__(
             master,
-            value,
             self.title,
             self.columns,
+            value,
             callback,
             *cb_args,
             **cb_kwargs
         )
 
 
-class FocusEditor(urwid.WidgetWrap):
+class FocusEditor(urwid.WidgetWrap, layoutwidget.LayoutWidget):
     """
         A specialised GridEditor that edits the current focused flow.
     """
@@ -442,27 +450,11 @@ class FocusEditor(urwid.WidgetWrap):
         self.master = master
         self.focus_changed()
 
-    def focus_changed(self):
-        if self.master.view.focus.flow:
-            self._w = BaseGridEditor(
-                self.master.view.focus.flow,
-                self.title,
-                self.columns,
-                self.get_data(self.master.view.focus.flow),
-                self.set_data_update,
-                self.master.view.focus.flow,
-            )
-        else:
-            self._w = urwid.Pile([])
-
     def call(self, v, name, *args, **kwargs):
         f = getattr(v, name, None)
         if f:
             f(*args, **kwargs)
 
-    def view_popping(self):
-        self.call(self._w, "view_popping")
-
     def get_data(self, flow):
         """
             Retrieve the data to edit from the current flow.
@@ -478,3 +470,22 @@ class FocusEditor(urwid.WidgetWrap):
     def set_data_update(self, vals, flow):
         self.set_data(vals, flow)
         signals.flow_change.send(self, flow = flow)
+
+    def key_responder(self):
+        return self._w
+
+    def view_popping(self):
+        self.call(self._w, "view_popping")
+
+    def focus_changed(self):
+        if self.master.view.focus.flow:
+            self._w = BaseGridEditor(
+                self.master.view.focus.flow,
+                self.title,
+                self.columns,
+                self.get_data(self.master.view.focus.flow),
+                self.set_data_update,
+                self.master.view.focus.flow,
+            )
+        else:
+            self._w = urwid.Pile([])
diff --git a/mitmproxy/tools/console/grideditor/editors.py b/mitmproxy/tools/console/grideditor/editors.py
index cd3c2c40..4a38eb78 100644
--- a/mitmproxy/tools/console/grideditor/editors.py
+++ b/mitmproxy/tools/console/grideditor/editors.py
@@ -1,11 +1,9 @@
-import re
 
 import urwid
 
 from mitmproxy import exceptions
-from mitmproxy import flowfilter
-from mitmproxy.addons import script
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 from mitmproxy.tools.console.grideditor import base
 from mitmproxy.tools.console.grideditor import col_text
 from mitmproxy.tools.console.grideditor import col_bytes
@@ -104,56 +102,6 @@ class RequestFormEditor(base.FocusEditor):
         flow.request.urlencoded_form = vals
 
 
-class SetHeadersEditor(base.GridEditor):
-    title = "Editing header set patterns"
-    columns = [
-        col_text.Column("Filter"),
-        col_text.Column("Header"),
-        col_text.Column("Value"),
-    ]
-
-    def is_error(self, col, val):
-        if col == 0:
-            if not flowfilter.parse(val):
-                return "Invalid filter specification"
-        return False
-
-    def make_help(self):
-        h = super().make_help()
-        text = [
-            urwid.Text([("text", "Special keys:\n")])
-        ]
-        keys = [
-            ("U", "add User-Agent header"),
-        ]
-        text.extend(
-            common.format_keyvals(keys, key="key", val="text", indent=4)
-        )
-        text.append(urwid.Text([("text", "\n")]))
-        text.extend(h)
-        return text
-
-    def set_user_agent(self, k):
-        ua = user_agents.get_by_shortcut(k)
-        if ua:
-            self.walker.add_value(
-                [
-                    ".*",
-                    b"User-Agent",
-                    ua[2].encode()
-                ]
-            )
-
-    def handle_key(self, key):
-        if key == "U":
-            signals.status_prompt_onekey.send(
-                prompt="Add User-Agent header:",
-                keys=[(i[0], i[1]) for i in user_agents.UASTRINGS],
-                callback=self.set_user_agent,
-            )
-            return True
-
-
 class PathEditor(base.FocusEditor):
     # TODO: Next row on enter?
 
@@ -175,38 +123,6 @@ class PathEditor(base.FocusEditor):
         flow.request.path_components = self.data_out(vals)
 
 
-class ScriptEditor(base.GridEditor):
-    title = "Editing scripts"
-    columns = [
-        col_text.Column("Command"),
-    ]
-
-    def is_error(self, col, val):
-        try:
-            script.parse_command(val)
-        except exceptions.OptionsError as e:
-            return str(e)
-
-
-class HostPatternEditor(base.GridEditor):
-    title = "Editing host patterns"
-    columns = [
-        col_text.Column("Regex (matched on hostname:port / ip:port)")
-    ]
-
-    def is_error(self, col, val):
-        try:
-            re.compile(val, re.IGNORECASE)
-        except re.error as e:
-            return "Invalid regex: %s" % str(e)
-
-    def data_in(self, data):
-        return [[i] for i in data]
-
-    def data_out(self, data):
-        return [i[0] for i in data]
-
-
 class CookieEditor(base.FocusEditor):
     title = "Edit Cookies"
     columns = [
@@ -273,7 +189,7 @@ class SetCookieEditor(base.FocusEditor):
         flow.response.cookies = self.data_out(vals)
 
 
-class OptionsEditor(base.GridEditor):
+class OptionsEditor(base.GridEditor, layoutwidget.LayoutWidget):
     title = None  # type: str
     columns = [
         col_text.Column("")
diff --git a/mitmproxy/tools/console/help.py b/mitmproxy/tools/console/help.py
index da8f701c..60ed4ec9 100644
--- a/mitmproxy/tools/console/help.py
+++ b/mitmproxy/tools/console/help.py
@@ -4,6 +4,7 @@ import urwid
 
 from mitmproxy import flowfilter
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 
 from mitmproxy import version
 
@@ -13,7 +14,7 @@ footer = [
 ]
 
 
-class HelpView(urwid.ListBox):
+class HelpView(urwid.ListBox, layoutwidget.LayoutWidget):
     title = "Help"
     keyctx = "help"
 
diff --git a/mitmproxy/tools/console/layoutwidget.py b/mitmproxy/tools/console/layoutwidget.py
new file mode 100644
index 00000000..689be3c0
--- /dev/null
+++ b/mitmproxy/tools/console/layoutwidget.py
@@ -0,0 +1,34 @@
+
+
+class LayoutWidget:
+    """
+        All top-level layout widgets and all widgets that may be set in an
+        overlay must comply with this API.
+    """
+    keyctx = ""
+
+    def key_responder(self):
+        """
+            Returns the object responding to key input. Usually self, but may be
+            a wrapped object.
+        """
+        return self
+
+    def focus_changed(self):
+        """
+            The view focus has changed. Layout objects should implement the API
+            rather than directly subscribing to events.
+        """
+        pass
+
+    def view_changed(self):
+        """
+            The view list has changed.
+        """
+        pass
+
+    def view_popping(self):
+        """
+            We are just about to pop a window off the stack, or exit an overlay.
+        """
+        pass
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 998d452d..6c1f3243 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -322,6 +322,40 @@ class ConsoleAddon:
                 "console.command flow.set @focus %s " % part
             )
 
+    def _grideditor(self):
+        gewidget = self.master.window.current("grideditor")
+        if not gewidget:
+            raise exceptions.CommandError("Not in a grideditor.")
+        return gewidget.key_responder()
+
+    @command.command("console.grideditor.add")
+    def grideditor_add(self) -> None:
+        """
+            Add a row after the cursor.
+        """
+        self._grideditor().cmd_add()
+
+    @command.command("console.grideditor.insert")
+    def grideditor_insert(self) -> None:
+        """
+            Insert a row before the cursor.
+        """
+        self._grideditor().cmd_insert()
+
+    @command.command("console.grideditor.next")
+    def grideditor_next(self) -> None:
+        """
+            Go to next cell.
+        """
+        self._grideditor().cmd_next()
+
+    @command.command("console.grideditor.delete")
+    def grideditor_delete(self) -> None:
+        """
+            Delete row
+        """
+        self._grideditor().cmd_delete()
+
     @command.command("console.flowview.mode.set")
     def flowview_mode_set(self) -> None:
         """
@@ -349,7 +383,7 @@ class ConsoleAddon:
         """
             Get the display mode for the current flow view.
         """
-        fv = self.master.window.any("flowview")
+        fv = self.master.window.current_window("flowview")
         if not fv:
             raise exceptions.CommandError("Not viewing a flow.")
         idx = fv.body.tab_offset
@@ -476,6 +510,11 @@ def default_keymap(km):
     km.add("D", "options.reset", ["options"])
     km.add("d", "console.options.reset.current", ["options"])
 
+    km.add("a", "console.grideditor.add", ["grideditor"])
+    km.add("A", "console.grideditor.insert", ["grideditor"])
+    km.add("tab", "console.grideditor.next", ["grideditor"])
+    km.add("d", "console.grideditor.delete", ["grideditor"])
+
 
 class ConsoleMaster(master.Master):
 
diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py
index 124a3f93..c364c308 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -7,6 +7,7 @@ from typing import Optional, Sequence
 from mitmproxy import exceptions
 from mitmproxy import optmanager
 from mitmproxy.tools.console import common
+from mitmproxy.tools.console import layoutwidget
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import overlay
 
@@ -263,7 +264,7 @@ class OptionHelp(urwid.Frame):
         self.set_body(self.widget(txt))
 
 
-class Options(urwid.Pile):
+class Options(urwid.Pile, layoutwidget.LayoutWidget):
     title = "Options"
     keyctx = "options"
 
diff --git a/mitmproxy/tools/console/overlay.py b/mitmproxy/tools/console/overlay.py
index abfb3909..dbf69f2e 100644
--- a/mitmproxy/tools/console/overlay.py
+++ b/mitmproxy/tools/console/overlay.py
@@ -5,10 +5,10 @@ import urwid
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import grideditor
+from mitmproxy.tools.console import layoutwidget
 
 
-class SimpleOverlay(urwid.Overlay):
-    keyctx = "overlay"
+class SimpleOverlay(urwid.Overlay, layoutwidget.LayoutWidget):
 
     def __init__(self, master, widget, parent, width, valign="middle"):
         self.widget = widget
@@ -22,14 +22,21 @@ class SimpleOverlay(urwid.Overlay):
             height="pack"
         )
 
-    def keypress(self, size, key):
-        key = super().keypress(size, key)
-        if key == "esc":
-            signals.pop_view_state.send(self)
-        if key == "?":
-            self.master.view_help(self.widget.make_help())
-        else:
-            return key
+    @property
+    def keyctx(self):
+        return getattr(self.widget, "keyctx")
+
+    def key_responder(self):
+        return self.widget.key_responder()
+
+    def focus_changed(self):
+        return self.widget.focus_changed()
+
+    def view_changed(self):
+        return self.widget.view_changed()
+
+    def view_popping(self):
+        return self.widget.view_popping()
 
 
 class Choice(urwid.WidgetWrap):
@@ -81,7 +88,9 @@ class ChooserListWalker(urwid.ListWalker):
         return self._get(pos, False), pos
 
 
-class Chooser(urwid.WidgetWrap):
+class Chooser(urwid.WidgetWrap, layoutwidget.LayoutWidget):
+    keyctx = "chooser"
+
     def __init__(self, master, title, choices, current, callback):
         self.master = master
         self.choices = choices
@@ -122,7 +131,9 @@ class Chooser(urwid.WidgetWrap):
         return text
 
 
-class OptionsOverlay(urwid.WidgetWrap):
+class OptionsOverlay(urwid.WidgetWrap, layoutwidget.LayoutWidget):
+    keyctx = "grideditor"
+
     def __init__(self, master, name, vals, vspace):
         """
             vspace: how much vertical space to keep clear
@@ -142,3 +153,9 @@ class OptionsOverlay(urwid.WidgetWrap):
 
     def make_help(self):
         return self.ge.make_help()
+
+    def key_responder(self):
+        return self.ge.key_responder()
+
+    def view_popping(self):
+        return self.ge.view_popping()
diff --git a/mitmproxy/tools/console/window.py b/mitmproxy/tools/console/window.py
index 964429df..51f36c78 100644
--- a/mitmproxy/tools/console/window.py
+++ b/mitmproxy/tools/console/window.py
@@ -45,16 +45,23 @@ class WindowStack:
         self.overlay = None
 
     def set_overlay(self, o, **kwargs):
-        self.overlay = overlay.SimpleOverlay(self, o, self.top(), o.width, **kwargs)
+        self.overlay = overlay.SimpleOverlay(
+            self, o, self.top_widget(), o.width, **kwargs,
+        )
 
-    @property
-    def topwin(self):
+    def top_window(self):
+        """
+            The current top window, ignoring overlays.
+        """
         return self.windows[self.stack[-1]]
 
-    def top(self):
+    def top_widget(self):
+        """
+            The current top widget - either a window or the active overlay.
+        """
         if self.overlay:
             return self.overlay
-        return self.topwin
+        return self.top_window()
 
     def push(self, wname):
         if self.stack[-1] == wname:
@@ -66,6 +73,7 @@ class WindowStack:
             Pop off the stack, return True if we're already at the top.
         """
         if self.overlay:
+            self.call("view_popping")
             self.overlay = None
         elif len(self.stack) > 1:
             self.call("view_popping")
@@ -74,9 +82,14 @@ class WindowStack:
             return True
 
     def call(self, name, *args, **kwargs):
-        f = getattr(self.topwin, name, None)
-        if f:
-            f(*args, **kwargs)
+        """
+            Call a function on both the top window, and the overlay if there is
+            one. If the widget has a key_responder, we call the function on the
+            responder instead.
+        """
+        getattr(self.top_window(), name)(*args, **kwargs)
+        if self.overlay:
+            getattr(self.overlay, name)(*args, **kwargs)
 
 
 class Window(urwid.Frame):
@@ -130,16 +143,16 @@ class Window(urwid.Frame):
 
         w = None
         if c == "single":
-            w = wrap(self.stacks[0].top(), 0)
+            w = wrap(self.stacks[0].top_widget(), 0)
         elif c == "vertical":
             w = urwid.Pile(
                 [
-                    wrap(s.top(), i) for i, s in enumerate(self.stacks)
+                    wrap(s.top_widget(), i) for i, s in enumerate(self.stacks)
                 ]
             )
         else:
             w = urwid.Columns(
-                [wrap(s.top(), i) for i, s in enumerate(self.stacks)],
+                [wrap(s.top_widget(), i) for i, s in enumerate(self.stacks)],
                 dividechars=1
             )
 
@@ -195,11 +208,18 @@ class Window(urwid.Frame):
 
     def current(self, keyctx):
         """
+            Returns the active widget, but only the current focus or overlay has
+            a matching key context.
+        """
+        t = self.focus_stack().top_widget()
+        if t.keyctx == keyctx:
+            return t
 
-            Returns the top window of the current stack, IF the current focus
-            has a matching key context.
+    def current_window(self, keyctx):
+        """
+            Returns the active window, ignoring overlays.
         """
-        t = self.focus_stack().topwin
+        t = self.focus_stack().top_window()
         if t.keyctx == keyctx:
             return t
 
@@ -207,7 +227,7 @@ class Window(urwid.Frame):
         """
             Returns the top window of either stack if they match the context.
         """
-        for t in [x.topwin for x in self.stacks]:
+        for t in [x.top_window() for x in self.stacks]:
             if t.keyctx == keyctx:
                 return t
 
@@ -245,7 +265,7 @@ class Window(urwid.Frame):
         if self.focus_part == "footer":
             return super().keypress(size, k)
         else:
-            fs = self.focus_stack().top()
+            fs = self.focus_stack().top_widget()
             k = fs.keypress(size, k)
             if k:
                 return self.master.keymap.handle(fs.keyctx, k)