Introduce a custom widget for command editing

The builtin urwid.Edit widget is not sufficiently flexible for what we want to
do.

5 files changed, 126 insertions, 11 deletions

diff --git a/mitmproxy/tools/console/commandeditor.py b/mitmproxy/tools/console/commandeditor.py
index 17d1506b..e57ddbb4 100644
--- a/mitmproxy/tools/console/commandeditor.py
+++ b/mitmproxy/tools/console/commandeditor.py
@@ -1,19 +1,10 @@
 import typing
-import urwid
 
 from mitmproxy import exceptions
 from mitmproxy import flow
 from mitmproxy.tools.console import signals
 
 
-class CommandEdit(urwid.Edit):
-    def __init__(self, partial):
-        urwid.Edit.__init__(self, ":", partial)
-
-    def keypress(self, size, key):
-        return urwid.Edit.keypress(self, size, key)
-
-
 class CommandExecutor:
     def __init__(self, master):
         self.master = master
diff --git a/mitmproxy/tools/console/commander/__init__.py b/mitmproxy/tools/console/commander/__init__.py
new file mode 100644
index 00000000..8b137891
--- /dev/null
+++ b/mitmproxy/tools/console/commander/__init__.py
@@ -0,0 +1 @@
+
diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
new file mode 100644
index 00000000..74855e4d
--- /dev/null
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -0,0 +1,85 @@
+import urwid
+from urwid.text_layout import calc_coords
+
+
+class CommandBuffer():
+    def __init__(self, start: str = ""):
+        self.buf = start
+        # This is the logical cursor position - the display cursor is one
+        # character further on. Cursor is always within the range [0:len(buffer)].
+        self._cursor = len(self.buf)
+
+    @property
+    def cursor(self):
+        return self._cursor
+
+    @cursor.setter
+    def cursor(self, x):
+        if x < 0:
+            self._cursor = 0
+        elif x > len(self.buf):
+            self._cursor = len(self.buf)
+        else:
+            self._cursor = x
+
+    def render(self):
+        return self.buf
+
+    def left(self):
+        self.cursor = self.cursor - 1
+
+    def right(self):
+        self.cursor = self.cursor + 1
+
+    def backspace(self):
+        if self.cursor == 0:
+            return
+        self.buf = self.buf[:self.cursor - 1] + self.buf[self.cursor:]
+        self.cursor = self.cursor - 1
+
+    def insert(self, k: str):
+        """
+            Inserts text at the cursor.
+        """
+        self.buf = self.buf = self.buf[:self.cursor] + k + self.buf[self.cursor:]
+        self.cursor += 1
+
+
+class CommandEdit(urwid.WidgetWrap):
+    leader = ": "
+
+    def __init__(self, text):
+        self.cbuf = CommandBuffer(text)
+        self._w = urwid.Text(self.leader)
+        self.update()
+
+    def keypress(self, size, key):
+        if key == "backspace":
+            self.cbuf.backspace()
+        elif key == "left":
+            self.cbuf.left()
+        elif key == "right":
+            self.cbuf.right()
+        elif len(key) == 1:
+            self.cbuf.insert(key)
+        self.update()
+
+    def update(self):
+        self._w.set_text([self.leader, self.cbuf.render()])
+
+    def render(self, size, focus=False):
+        (maxcol,) = size
+        canv = self._w.render((maxcol,))
+        canv = urwid.CompositeCanvas(canv)
+        canv.cursor = self.get_cursor_coords((maxcol,))
+        return canv
+
+    def get_cursor_coords(self, size):
+        p = self.cbuf.cursor + len(self.leader)
+        trans = self._w.get_line_translation(size[0])
+        x, y = calc_coords(self._w.get_text()[0], trans, p)
+        return x, y
+
+    def get_value(self):
+        return self.cbuf.buf
+
diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py
index 795b3d8a..a59fc92e 100644
--- a/mitmproxy/tools/console/statusbar.py
+++ b/mitmproxy/tools/console/statusbar.py
@@ -6,6 +6,7 @@ from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import commandeditor
 import mitmproxy.tools.console.master # noqa
+from mitmproxy.tools.console.commander import commander
 
 
 class PromptPath:
@@ -66,7 +67,7 @@ class ActionBar(urwid.WidgetWrap):
 
     def sig_prompt_command(self, sender, partial=""):
         signals.focus.send(self, section="footer")
-        self._w = commandeditor.CommandEdit(partial)
+        self._w = commander.CommandEdit(partial)
         self.prompting = commandeditor.CommandExecutor(self.master)
 
     def sig_prompt_onekey(self, sender, prompt, keys, callback, args=()):
@@ -100,7 +101,7 @@ class ActionBar(urwid.WidgetWrap):
                 elif k in self.onekey:
                     self.prompt_execute(k)
             elif k == "enter":
-                self.prompt_execute(self._w.get_edit_text())
+                self.prompt_execute(self._w.get_value())
             else:
                 if common.is_keypress(k):
                     self._w.keypress(size, k)
diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py
new file mode 100644
index 00000000..b1f23df4
--- /dev/null
+++ b/test/mitmproxy/tools/console/test_commander.py
@@ -0,0 +1,37 @@
+
+from mitmproxy.tools.console.commander import commander
+
+
+class TestCommandBuffer:
+
+    def test_backspace(self):
+        tests = [
+            [("", 0), ("", 0)],
+            [("1", 0), ("1", 0)],
+            [("1", 1), ("", 0)],
+            [("123", 3), ("12", 2)],
+            [("123", 2), ("13", 1)],
+            [("123", 0), ("123", 0)],
+        ]
+        for start, output in tests:
+            cb = commander.CommandBuffer()
+            cb.buf, cb.cursor = start[0], start[1]
+            cb.backspace()
+            assert cb.buf == output[0]
+            assert cb.cursor == output[1]
+
+    def test_insert(self):
+        tests = [
+            [("", 0), ("x", 1)],
+            [("a", 0), ("xa", 1)],
+            [("xa", 2), ("xax", 3)],
+        ]
+        for start, output in tests:
+            cb = commander.CommandBuffer()
+            cb.buf, cb.cursor = start[0], start[1]
+            cb.insert("x")
+            assert cb.buf == output[0]
+            assert cb.cursor == output[1]
+
+
+