Merge pull request #2394 from cortesi/moarconsole

Misc console-related improvements

mitmproxy/addons/core.py

@@ -10,15 +10,17 @@ from mitmproxy.net.http import status_codes
 
 class Core:
     @command.command("set")
-    def set(self, spec: str) -> None:
+    def set(self, *spec: str) -> None:
         """
             Set an option of the form "key[=value]". When the value is omitted,
             booleans are set to true, strings and integers are set to None (if
             permitted), and sequences are emptied. Boolean values can be true,
-            false or toggle.
+            false or toggle. If multiple specs are passed, they are joined
+            into one separated by spaces.
         """
+        strspec = " ".join(spec)
         try:
-            ctx.options.set(spec)
+            ctx.options.set(strspec)
         except exceptions.OptionsError as e:
             raise exceptions.CommandError(e) from e
 

mitmproxy/addons/view.py

@@ -389,6 +389,8 @@ class View(collections.Sequence):
                     self.sig_view_remove.send(self, flow=f)
                 del self._store[f.id]
                 self.sig_store_remove.send(self, flow=f)
+        if len(flows) > 1:
+            ctx.log.alert("Removed %s flows" % len(flows))
 
     @command.command("view.resolve")
     def resolve(self, spec: str) -> typing.Sequence[mitmproxy.flow.Flow]:

mitmproxy/command.py

@@ -59,7 +59,7 @@ class Command:
     def paramnames(self) -> typing.Sequence[str]:
         v = [typename(i, False) for i in self.paramtypes]
         if self.has_positional:
-            v[-1] = "*" + v[-1][1:-1]
+            v[-1] = "*" + v[-1]
         return v
 
     def retname(self) -> str:
@@ -92,7 +92,11 @@ class Command:
                 pargs.append(parsearg(self.manager, args[i], self.paramtypes[i]))
 
         if remainder:
-            if typecheck.check_command_type(remainder, self.paramtypes[-1]):
+            chk = typecheck.check_command_type(
+                remainder,
+                typing.Sequence[self.paramtypes[-1]]  # type: ignore
+            )
+            if chk:
                 pargs.extend(remainder)
             else:
                 raise exceptions.CommandError("Invalid value type.")

mitmproxy/options.py

@@ -394,7 +394,7 @@ class Options(optmanager.OptManager):
             "Focus follows new flows."
         )
         self.add_option(
-            "console_palette", str, "dark",
+            "console_palette", str, "solarized_dark",
             "Color palette.",
             choices=sorted(console_palettes),
         )

mitmproxy/optmanager.py

@@ -432,6 +432,8 @@ def parse(text):
             raise exceptions.OptionsError("Could not parse options.")
     if isinstance(data, str):
         raise exceptions.OptionsError("Config error - no keys found.")
+    elif data is None:
+        return {}
     return data
 
 

mitmproxy/tools/console/commands.py

@@ -6,16 +6,6 @@ from mitmproxy.tools.console import signals
 
 HELP_HEIGHT = 5
 
-
-def fcol(s, width, attr):
-    s = str(s)
-    return (
-        "fixed",
-        width,
-        urwid.Text((attr, s))
-    )
-
-
 command_focus_change = blinker.Signal()
 
 

mitmproxy/tools/console/defaultkeys.py

@@ -3,6 +3,7 @@ def map(km):
     km.add(":", "console.command ''", ["global"], "Command prompt")
     km.add("?", "console.view.help", ["global"], "View help")
     km.add("C", "console.view.commands", ["global"], "View commands")
+    km.add("K", "console.view.keybindings", ["global"], "View key bindings")
     km.add("O", "console.view.options", ["global"], "View options")
     km.add("E", "console.view.eventlog", ["global"], "View event log")
     km.add("Q", "console.exit", ["global"], "Exit immediately")
@@ -36,8 +37,10 @@ def map(km):
     km.add("D", "view.duplicate @focus", ["flowlist", "flowview"], "Duplicate flow")
     km.add(
         "e",
-        "console.choose.cmd Format export.formats "
-        "console.command export.file {choice} @focus ''",
+        """
+        console.choose.cmd Format export.formats
+        console.command export.file {choice} @focus ''
+        """,
         ["flowlist", "flowview"],
         "Export this flow to file"
     )
@@ -60,8 +63,10 @@ def map(km):
     )
     km.add(
         "o",
-        "console.choose.cmd Order view.order.options "
-        "set console_order={choice}",
+        """
+        console.choose.cmd Order view.order.options
+        set console_order={choice}
+        """,
         ["flowlist"],
         "Set flow list order"
     )
@@ -83,8 +88,10 @@ def map(km):
 
     km.add(
         "e",
-        "console.choose.cmd Part console.edit.focus.options "
-        "console.edit.focus {choice}",
+        """
+        console.choose.cmd Part console.edit.focus.options
+        console.edit.focus {choice}
+        """,
         ["flowview"],
         "Edit a flow component"
     )
@@ -99,8 +106,10 @@ def map(km):
 
     km.add(
         "v",
-        "console.choose \"View Part\" request,response "
-        "console.bodyview @focus {choice}",
+        """
+        console.choose "View Part" request,response
+        console.bodyview @focus {choice}
+        """,
         ["flowview"],
         "View flow body in an external viewer"
     )
@@ -108,8 +117,10 @@ def map(km):
     km.add("m", "console.flowview.mode.set", ["flowview"], "Set flow view mode")
     km.add(
         "z",
-        "console.choose \"Part\" request,response "
-        "flow.encode.toggle @focus {choice}",
+        """
+        console.choose "Part" request,response
+        flow.encode.toggle @focus {choice}
+        """,
         ["flowview"],
         "Encode/decode flow body"
     )

mitmproxy/tools/console/flowlist.py

@@ -30,7 +30,7 @@ class FlowItem(urwid.WidgetWrap):
                 self.master.commands.call("console.view.flow @focus")
                 return True
 
-    def keypress(self, xxx_todo_changeme, key):
+    def keypress(self, size, key):
         return key
 
 

mitmproxy/tools/console/keybindings.py

@@ -0,0 +1,146 @@
+import urwid
+import blinker
+import textwrap
+from mitmproxy.tools.console import layoutwidget
+
+HELP_HEIGHT = 5
+
+
+keybinding_focus_change = blinker.Signal()
+
+
+class KeyItem(urwid.WidgetWrap):
+    def __init__(self, walker, binding, focused):
+        self.walker, self.binding, self.focused = walker, binding, focused
+        super().__init__(None)
+        self._w = self.get_widget()
+
+    def get_widget(self):
+        cmd = textwrap.dedent(self.binding.command).strip()
+        parts = [
+            (4, urwid.Text([("focus", ">> " if self.focused else "   ")])),
+            (10, urwid.Text([("title", self.binding.key)])),
+            (12, urwid.Text([("highlight", "\n".join(self.binding.contexts))])),
+            urwid.Text([("text", cmd)]),
+        ]
+        return urwid.Columns(parts)
+
+    def get_edit_text(self):
+        return self._w[1].get_edit_text()
+
+    def selectable(self):
+        return True
+
+    def keypress(self, size, key):
+        return key
+
+
+class KeyListWalker(urwid.ListWalker):
+    def __init__(self, master):
+        self.master = master
+
+        self.index = 0
+        self.focusobj = None
+        self.bindings = list(master.keymap.list("all"))
+        self.set_focus(0)
+
+    def get_edit_text(self):
+        return self.focus_obj.get_edit_text()
+
+    def _get(self, pos):
+        binding = self.bindings[pos]
+        return KeyItem(self, binding, pos == self.index)
+
+    def get_focus(self):
+        return self.focus_obj, self.index
+
+    def set_focus(self, index):
+        binding = self.bindings[index]
+        self.index = index
+        self.focus_obj = self._get(self.index)
+        keybinding_focus_change.send(binding.help or "")
+
+    def get_next(self, pos):
+        if pos >= len(self.bindings) - 1:
+            return None, None
+        pos = pos + 1
+        return self._get(pos), pos
+
+    def get_prev(self, pos):
+        pos = pos - 1
+        if pos < 0:
+            return None, None
+        return self._get(pos), pos
+
+
+class KeyList(urwid.ListBox):
+    def __init__(self, master):
+        self.master = master
+        self.walker = KeyListWalker(master)
+        super().__init__(self.walker)
+
+    def keypress(self, size, key):
+        if key == "m_select":
+            foc, idx = self.get_focus()
+            # Act here
+        elif key == "m_start":
+            self.set_focus(0)
+            self.walker._modified()
+        elif key == "m_end":
+            self.set_focus(len(self.walker.bindings) - 1)
+            self.walker._modified()
+        return super().keypress(size, key)
+
+
+class KeyHelp(urwid.Frame):
+    def __init__(self, master):
+        self.master = master
+        super().__init__(self.widget(""))
+        self.set_active(False)
+        keybinding_focus_change.connect(self.sig_mod)
+
+    def set_active(self, val):
+        h = urwid.Text("Key Binding Help")
+        style = "heading" if val else "heading_inactive"
+        self.header = urwid.AttrWrap(h, style)
+
+    def widget(self, txt):
+        cols, _ = self.master.ui.get_cols_rows()
+        return urwid.ListBox(
+            [urwid.Text(i) for i in textwrap.wrap(txt, cols)]
+        )
+
+    def sig_mod(self, txt):
+        self.set_body(self.widget(txt))
+
+
+class KeyBindings(urwid.Pile, layoutwidget.LayoutWidget):
+    title = "Key Bindings"
+    keyctx = "keybindings"
+
+    def __init__(self, master):
+        oh = KeyHelp(master)
+        super().__init__(
+            [
+                KeyList(master),
+                (HELP_HEIGHT, oh),
+            ]
+        )
+        self.master = master
+
+    def keypress(self, size, key):
+        if key == "m_next":
+            self.focus_position = (
+                self.focus_position + 1
+            ) % len(self.widget_list)
+            self.widget_list[1].set_active(self.focus_position == 1)
+            key = None
+
+        # This is essentially a copypasta from urwid.Pile's keypress handler.
+        # So much for "closed for modification, but open for extension".
+        item_rows = None
+        if len(size) == 2:
+            item_rows = self.get_item_rows(size, focus = True)
+        i = self.widget_list.index(self.focus_item)
+        tsize = self.get_item_size(size, i, True, item_rows)
+        return self.focus_item.keypress(tsize, key)

mitmproxy/tools/console/keymap.py

@@ -54,7 +54,7 @@ class Keymap:
         return None
 
     def list(self, context: str) -> typing.Sequence[Binding]:
-        b = [b for b in self.bindings if context in b.contexts]
+        b = [b for b in self.bindings if context in b.contexts or context == "all"]
         b.sort(key=lambda x: x.key)
         return b
 

mitmproxy/tools/console/master.py

@@ -189,7 +189,7 @@ class ConsoleAddon:
 
     @command.command("console.choose")
     def console_choose(
-        self, prompt: str, choices: typing.Sequence[str], *cmd: typing.Sequence[str]
+        self, prompt: str, choices: typing.Sequence[str], *cmd: str
     ) -> None:
         """
             Prompt the user to choose from a specified list of strings, then
@@ -211,7 +211,7 @@ class ConsoleAddon:
 
     @command.command("console.choose.cmd")
     def console_choose_cmd(
-        self, prompt: str, choicecmd: str, *cmd: typing.Sequence[str]
+        self, prompt: str, choicecmd: str, *cmd: str
     ) -> None:
         """
             Prompt the user to choose from a list of strings returned by a
@@ -234,11 +234,16 @@ class ConsoleAddon:
         )
 
     @command.command("console.command")
-    def console_command(self, *partial: typing.Sequence[str]) -> None:
+    def console_command(self, *partial: str) -> None:
         """
         Prompt the user to edit a command with a (possilby empty) starting value.
         """
-        signals.status_prompt_command.send(partial=" ".join(partial) + " ")  # type: ignore
+        signals.status_prompt_command.send(partial=" ".join(partial))  # type: ignore
+
+    @command.command("console.view.keybindings")
+    def view_keybindings(self) -> None:
+        """View the commands list."""
+        self.master.switch_view("keybindings")
 
     @command.command("console.view.commands")
     def view_commands(self) -> None:

mitmproxy/tools/console/window.py

@@ -4,6 +4,7 @@ from mitmproxy.tools.console import statusbar
 from mitmproxy.tools.console import flowlist
 from mitmproxy.tools.console import flowview
 from mitmproxy.tools.console import commands
+from mitmproxy.tools.console import keybindings
 from mitmproxy.tools.console import options
 from mitmproxy.tools.console import overlay
 from mitmproxy.tools.console import help
@@ -29,6 +30,7 @@ class WindowStack:
             flowlist = flowlist.FlowListBox(master),
             flowview = flowview.FlowView(master),
             commands = commands.Commands(master),
+            keybindings = keybindings.KeyBindings(master),
             options = options.Options(master),
             help = help.HelpView(master),
             eventlog = eventlog.EventLog(master),

test/mitmproxy/addons/test_view.py

@@ -262,6 +262,16 @@ def test_duplicate():
         assert v.focus.index == 2
 
 
+def test_remove():
+    v = view.View()
+    with taddons.context():
+        f = [tflow.tflow(), tflow.tflow()]
+        v.add(f)
+        assert len(v) == 2
+        v.remove(f)
+        assert len(v) == 0
+
+
 def test_setgetval():
     v = view.View()
     with taddons.context():

test/mitmproxy/test_command.py

@@ -22,7 +22,7 @@ class TAddon:
     def empty(self) -> None:
         pass
 
-    def varargs(self, one: str, *var: typing.Sequence[str]) -> typing.Sequence[str]:
+    def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
         return list(var)
 
 

test/mitmproxy/test_optmanager.py

@@ -257,6 +257,10 @@ def test_serialize():
     with pytest.raises(Exception, match="Config error"):
         optmanager.load(o2, t)
 
+    t = "# a comment"
+    optmanager.load(o2, t)
+    assert optmanager.load(o2, "foobar: '123'") == {"foobar": "123"}
+
     t = ""
     optmanager.load(o2, t)
     assert optmanager.load(o2, "foobar: '123'") == {"foobar": "123"}