From 04e19f91716b9de6ec26df1478146eaedd47a329 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 14 Dec 2017 12:24:46 +1300 Subject: [PATCH] Introduce a custom widget for command editing The builtin urwid.Edit widget is not sufficiently flexible for what we want to do. --- mitmproxy/tools/console/commandeditor.py | 9 -- mitmproxy/tools/console/commander/__init__.py | 1 + .../tools/console/commander/commander.py | 85 +++++++++++++++++++ mitmproxy/tools/console/statusbar.py | 5 +- .../mitmproxy/tools/console/test_commander.py | 37 ++++++++ 5 files changed, 126 insertions(+), 11 deletions(-) create mode 100644 mitmproxy/tools/console/commander/__init__.py create mode 100644 mitmproxy/tools/console/commander/commander.py create mode 100644 test/mitmproxy/tools/console/test_commander.py diff --git a/mitmproxy/tools/console/commandeditor.py b/mitmproxy/tools/console/commandeditor.py index 17d1506bd..e57ddbb4d 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 000000000..8b1378917 --- /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 000000000..74855e4de --- /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 795b3d8a9..a59fc92e1 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 000000000..b1f23df42 --- /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] + + +