Merge pull request #2716 from cortesi/cmdfixes

various command-related fixes

8 files changed, 49 insertions, 3 deletions

diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py
index b90df549..f4b560e8 100644
--- a/mitmproxy/addons/cut.py
+++ b/mitmproxy/addons/cut.py
@@ -36,6 +36,8 @@ def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
             if spec == "host" and is_addr(current):
                 return str(current[0])
             elif spec.startswith("header["):
+                if not current:
+                    return ""
                 return current.headers.get(headername(spec), "")
             elif isinstance(part, bytes):
                 return part
diff --git a/mitmproxy/tools/console/commandexecutor.py b/mitmproxy/tools/console/commandexecutor.py
index 829daee1..26f92238 100644
--- a/mitmproxy/tools/console/commandexecutor.py
+++ b/mitmproxy/tools/console/commandexecutor.py
@@ -23,6 +23,10 @@ class CommandExecutor:
                         signals.status_message.send(
                             message="Command returned %s flows" % len(ret)
                         )
+                    elif type(ret) == flow.Flow:
+                        signals.status_message.send(
+                            message="Command returned 1 flow"
+                        )
                     else:
                         self.master.overlay(
                             overlay.DataViewerOverlay(
diff --git a/mitmproxy/tools/console/commands.py b/mitmproxy/tools/console/commands.py
index 20efcee3..1183ee9d 100644
--- a/mitmproxy/tools/console/commands.py
+++ b/mitmproxy/tools/console/commands.py
@@ -124,7 +124,7 @@ class CommandHelp(urwid.Frame):
 
 
 class Commands(urwid.Pile, layoutwidget.LayoutWidget):
-    title = "Commands"
+    title = "Command Reference"
     keyctx = "commands"
 
     def __init__(self, master):
diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py
index ce3e57bd..20d54bc6 100644
--- a/mitmproxy/tools/console/consoleaddons.py
+++ b/mitmproxy/tools/console/consoleaddons.py
@@ -322,6 +322,7 @@ class ConsoleAddon:
         signals.pop_view_state.send(self)
 
     @command.command("console.bodyview")
+    @command.argument("part", type=mitmproxy.types.Choice("console.bodyview.options"))
     def bodyview(self, f: flow.Flow, part: str) -> None:
         """
             Spawn an external viewer for a flow request or response body based
@@ -338,6 +339,13 @@ class ConsoleAddon:
             raise exceptions.CommandError("No content to view.")
         self.master.spawn_external_viewer(content, t)
 
+    @command.command("console.bodyview.options")
+    def bodyview_options(self) -> typing.Sequence[str]:
+        """
+            Possible parts for console.bodyview.
+        """
+        return ["request", "response"]
+
     @command.command("console.edit.focus.options")
     def edit_focus_options(self) -> typing.Sequence[str]:
         """
diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py
index b268906c..fbb569a4 100644
--- a/mitmproxy/tools/console/keymap.py
+++ b/mitmproxy/tools/console/keymap.py
@@ -17,6 +17,13 @@ Contexts = {
 }
 
 
+navkeys = [
+    "m_start", "m_end", "m_next", "m_select",
+    "up", "down", "page_up", "page_down",
+    "left", "right"
+]
+
+
 class Binding:
     def __init__(self, key, command, contexts, help):
         self.key, self.command, self.contexts = key, command, sorted(contexts)
@@ -122,3 +129,13 @@ class Keymap:
         if b:
             return self.executor(b.command)
         return key
+
+    def handle_only(self, context: str, key: str) -> typing.Optional[str]:
+        """
+            Like handle, but ignores global bindings. Returns the key if it has
+            not been handled, or None.
+        """
+        b = self.get(context, key)
+        if b:
+            return self.executor(b.command)
+        return key
diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py
index 4d55aeec..54772cf0 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -117,6 +117,7 @@ class OptionListWalker(urwid.ListWalker):
     def stop_editing(self):
         self.editing = False
         self.focus_obj = self._get(self.index, False)
+        self.set_focus(self.index)
         self._modified()
 
     def get_edit_text(self):
diff --git a/mitmproxy/tools/console/overlay.py b/mitmproxy/tools/console/overlay.py
index f97f23f9..55acbfdd 100644
--- a/mitmproxy/tools/console/overlay.py
+++ b/mitmproxy/tools/console/overlay.py
@@ -5,6 +5,7 @@ import urwid
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import grideditor
 from mitmproxy.tools.console import layoutwidget
+from mitmproxy.tools.console import keymap
 
 
 class SimpleOverlay(urwid.Overlay, layoutwidget.LayoutWidget):
@@ -114,13 +115,21 @@ class Chooser(urwid.WidgetWrap, layoutwidget.LayoutWidget):
         return True
 
     def keypress(self, size, key):
-        key = self.master.keymap.handle("chooser", key)
+        key = self.master.keymap.handle_only("chooser", key)
         if key == "m_select":
             self.callback(self.choices[self.walker.index])
             signals.pop_view_state.send(self)
+            return
         elif key == "esc":
             signals.pop_view_state.send(self)
-        return super().keypress(size, key)
+            return
+
+        binding = self.master.keymap.get("global", key)
+        # This is extremely awkward. We need a better way to match nav keys only.
+        if binding and binding.command.startswith("console.nav"):
+            self.master.keymap.handle("global", key)
+        elif key in keymap.navkeys:
+            return super().keypress(size, key)
 
 
 class OptionsOverlay(urwid.WidgetWrap, layoutwidget.LayoutWidget):
diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py
index 0a523fff..71e699db 100644
--- a/test/mitmproxy/addons/test_cut.py
+++ b/test/mitmproxy/addons/test_cut.py
@@ -135,6 +135,11 @@ def test_cut():
         with pytest.raises(exceptions.CommandError):
             assert c.cut(tflows, ["__dict__"]) == [[""]]
 
+    with taddons.context():
+        tflows = [tflow.tflow(resp=False)]
+        assert c.cut(tflows, ["response.reason"]) == [[""]]
+        assert c.cut(tflows, ["response.header[key]"]) == [[""]]
+
     c = cut.Cut()
     with taddons.context():
         tflows = [tflow.ttcpflow()]