Merge pull request #3251 from kajojify/command_history

Command history

3 files changed, 116 insertions, 2 deletions

diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index df3eaa5a..e8550f86 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -1,5 +1,7 @@
 import abc
+import copy
 import typing
+import collections
 
 import urwid
 from urwid.text_layout import calc_coords
@@ -156,13 +158,53 @@ class CommandBuffer:
         self.completion = None
 
 
+class CommandHistory:
+    def __init__(self, master: mitmproxy.master.Master, size: int=30) -> None:
+        self.saved_commands: collections.deque = collections.deque(
+            [CommandBuffer(master, "")],
+            maxlen=size
+        )
+        self.index: int = 0
+
+    @property
+    def last_index(self):
+        return len(self.saved_commands) - 1
+
+    def get_next(self) -> typing.Optional[CommandBuffer]:
+        if self.index < self.last_index:
+            self.index = self.index + 1
+            return self.saved_commands[self.index]
+        return None
+
+    def get_prev(self) -> typing.Optional[CommandBuffer]:
+        if self.index > 0:
+            self.index = self.index - 1
+            return self.saved_commands[self.index]
+        return None
+
+    def add_command(self, command: CommandBuffer, execution: bool=False) -> None:
+        if self.index == self.last_index or execution:
+            last_item = self.saved_commands[-1]
+            last_item_empty = not last_item.text
+            if last_item.text == command.text or (last_item_empty and execution):
+                self.saved_commands[-1] = copy.copy(command)
+            else:
+                self.saved_commands.append(command)
+                if not execution and self.index < self.last_index:
+                    self.index += 1
+            if execution:
+                self.index = self.last_index
+
+
 class CommandEdit(urwid.WidgetWrap):
     leader = ": "
 
-    def __init__(self, master: mitmproxy.master.Master, text: str) -> None:
+    def __init__(self, master: mitmproxy.master.Master,
+                 text: str, history: CommandHistory) -> None:
         super().__init__(urwid.Text(self.leader))
         self.master = master
         self.cbuf = CommandBuffer(master, text)
+        self.history = history
         self.update()
 
     def keypress(self, size, key):
@@ -172,6 +214,11 @@ class CommandEdit(urwid.WidgetWrap):
             self.cbuf.left()
         elif key == "right":
             self.cbuf.right()
+        elif key == "up":
+            self.history.add_command(self.cbuf)
+            self.cbuf = self.history.get_prev() or self.cbuf
+        elif key == "down":
+            self.cbuf = self.history.get_next() or self.cbuf
         elif key == "tab":
             self.cbuf.cycle_completion()
         elif len(key) == 1:
diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py
index 215cf500..2d32f487 100644
--- a/mitmproxy/tools/console/statusbar.py
+++ b/mitmproxy/tools/console/statusbar.py
@@ -42,6 +42,8 @@ class ActionBar(urwid.WidgetWrap):
         signals.status_prompt_onekey.connect(self.sig_prompt_onekey)
         signals.status_prompt_command.connect(self.sig_prompt_command)
 
+        self.command_history = commander.CommandHistory(master)
+
         self.prompting = None
 
         self.onekey = False
@@ -98,7 +100,8 @@ class ActionBar(urwid.WidgetWrap):
 
     def sig_prompt_command(self, sender, partial=""):
         signals.focus.send(self, section="footer")
-        self._w = commander.CommandEdit(self.master, partial)
+        self._w = commander.CommandEdit(self.master, partial,
+                                        self.command_history)
         self.prompting = commandexecutor.CommandExecutor(self.master)
 
     def sig_prompt_onekey(self, sender, prompt, keys, callback, args=()):
@@ -125,6 +128,7 @@ class ActionBar(urwid.WidgetWrap):
     def keypress(self, size, k):
         if self.prompting:
             if k == "esc":
+                self.command_history.index = self.command_history.last_index
                 self.prompt_done()
             elif self.onekey:
                 if k == "enter":
@@ -132,6 +136,7 @@ class ActionBar(urwid.WidgetWrap):
                 elif k in self.onekey:
                     self.prompt_execute(k)
             elif k == "enter":
+                self.command_history.add_command(self._w.cbuf, True)
                 self.prompt_execute(self._w.get_edit_text())
             else:
                 if common.is_keypress(k):
diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py
index 2a96995d..b5e226fe 100644
--- a/test/mitmproxy/tools/console/test_commander.py
+++ b/test/mitmproxy/tools/console/test_commander.py
@@ -28,6 +28,68 @@ class TestListCompleter:
                 assert c.cycle() == expected
 
 
+class TestCommandHistory:
+    def fill_history(self, commands):
+        with taddons.context() as tctx:
+            history = commander.CommandHistory(tctx.master, size=3)
+            for c in commands:
+                cbuf = commander.CommandBuffer(tctx.master, c)
+                history.add_command(cbuf)
+        return history, tctx.master
+
+    def test_add_command(self):
+        commands = ["command1", "command2"]
+        history, tctx_master = self.fill_history(commands)
+
+        saved_commands = [buf.text for buf in history.saved_commands]
+        assert saved_commands == [""] + commands
+
+        # The history size is only 3. So, we forget the first
+        # one command, when adding fourth command
+        cbuf = commander.CommandBuffer(tctx_master, "command3")
+        history.add_command(cbuf)
+        saved_commands = [buf.text for buf in history.saved_commands]
+        assert saved_commands == commands + ["command3"]
+
+        # Commands with the same text are not repeated in the history one by one
+        history.add_command(cbuf)
+        saved_commands = [buf.text for buf in history.saved_commands]
+        assert saved_commands == commands + ["command3"]
+
+        # adding command in execution mode sets index at the beginning of the history
+        # and replace the last command buffer if it is empty or has the same text
+        cbuf = commander.CommandBuffer(tctx_master, "")
+        history.add_command(cbuf)
+        history.index = 0
+        cbuf = commander.CommandBuffer(tctx_master, "command4")
+        history.add_command(cbuf, True)
+        assert history.index == history.last_index
+        saved_commands = [buf.text for buf in history.saved_commands]
+        assert saved_commands == ["command2", "command3", "command4"]
+
+    def test_get_next(self):
+        commands = ["command1", "command2"]
+        history, tctx_master = self.fill_history(commands)
+
+        history.index = -1
+        expected_items = ["", "command1", "command2"]
+        for i in range(3):
+            assert history.get_next().text == expected_items[i]
+        # We are at the last item of the history
+        assert history.get_next() is None
+
+    def test_get_prev(self):
+        commands = ["command1", "command2"]
+        history, tctx_master = self.fill_history(commands)
+
+        expected_items = ["command2", "command1", ""]
+        history.index = history.last_index + 1
+        for i in range(3):
+            assert history.get_prev().text == expected_items[i]
+        # We are at the first item of the history
+        assert history.get_prev() is None
+
+
 class TestCommandBuffer:
 
     def test_backspace(self):