From 16b55f9476373347a3c2553e070497b383288360 Mon Sep 17 00:00:00 2001
From: Henrique <typoon@gmail.com>
Date: Fri, 22 Nov 2019 10:00:17 -0500
Subject: Implemented feature to save command history to a file. This allows
 users to reuse their commands the next time they open mitmproxy

---
 mitmproxy/tools/console/commander/commander.py | 46 +++++++++++++++++++++++---
 test/mitmproxy/tools/console/test_commander.py |  7 ++--
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index d751422b..d661d530 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -2,6 +2,7 @@ import abc
 import collections
 import copy
 import typing
+import os
 
 import urwid
 from urwid.text_layout import calc_coords
@@ -147,29 +148,55 @@ class CommandBuffer:
         self.completion = None
 
 
+# TODO: This class should be a Singleton
 class CommandHistory:
-    def __init__(self, master: mitmproxy.master.Master, size: int = 30) -> None:
+    def __init__(self, master: mitmproxy.master.Master, size: int = 300) -> None:
         self.saved_commands: collections.deque = collections.deque(
             [CommandBuffer(master, "")],
             maxlen=size
         )
         self.index: int = 0
+        self.size: int = size
+        self.master: mitmproxy.master.Master = master
+
+        _command_history_path = os.path.join(os.path.expanduser(mitmproxy.options.CONF_DIR), 'command_history')
+        if os.path.exists(_command_history_path):
+            with open(_command_history_path, 'r') as f:
+                for l in f.readlines():
+                    cbuf = CommandBuffer(master, l.strip())
+                    self.add_command(cbuf)
+                f.close()
+
+        self.command_history_file = open(_command_history_path, 'w')
 
     @property
     def last_index(self):
         return len(self.saved_commands) - 1
 
+    def clear_history(self):
+        """
+        Needed for test suite.
+        TODO: Maybe create a command to clear the history?
+        """
+        self.saved_commands: collections.deque = collections.deque(
+            [CommandBuffer(self.master, "")],
+            maxlen=self.size
+        )
+
+        self.index = 0
+        self.command_history_file.truncate(0)
+        self.command_history_file.seek(0)
+        self.command_history_file.flush()
+
     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
+        return self.saved_commands[self.index]
 
     def get_prev(self) -> typing.Optional[CommandBuffer]:
         if self.index > 0:
             self.index = self.index - 1
-            return self.saved_commands[self.index]
-        return None
+        return self.saved_commands[self.index]
 
     def add_command(self, command: CommandBuffer, execution: bool = False) -> None:
         if self.index == self.last_index or execution:
@@ -184,6 +211,15 @@ class CommandHistory:
             if execution:
                 self.index = self.last_index
 
+            # This prevents the constructor from trying to overwrite the file
+            # that it is currently reading
+            if hasattr(self, 'command_history_file'):
+                _history_str = "\n".join([c.text for c in self.saved_commands])
+                self.command_history_file.truncate(0)
+                self.command_history_file.seek(0)
+                self.command_history_file.write(_history_str)
+                self.command_history_file.flush()
+
 
 class CommandEdit(urwid.WidgetWrap):
     leader = ": "
diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py
index a77be043..5e11166e 100644
--- a/test/mitmproxy/tools/console/test_commander.py
+++ b/test/mitmproxy/tools/console/test_commander.py
@@ -97,6 +97,7 @@ class TestCommandEdit:
     def test_up_and_down(self):
         with taddons.context() as tctx:
             history = commander.CommandHistory(tctx.master, size=3)
+            history.clear_history()
             edit = commander.CommandEdit(tctx.master, '', history)
 
             buf = commander.CommandBuffer(tctx.master, 'cmd1')
@@ -112,6 +113,7 @@ class TestCommandEdit:
             assert edit.get_edit_text() == 'cmd1'
 
             history = commander.CommandHistory(tctx.master, size=5)
+            history.clear_history()
             edit = commander.CommandEdit(tctx.master, '', history)
             edit.keypress(1, 'a')
             edit.keypress(1, 'b')
@@ -139,6 +141,7 @@ class TestCommandHistory:
     def fill_history(self, commands):
         with taddons.context() as tctx:
             history = commander.CommandHistory(tctx.master, size=3)
+            history.clear_history()
             for c in commands:
                 cbuf = commander.CommandBuffer(tctx.master, c)
                 history.add_command(cbuf)
@@ -183,7 +186,7 @@ class TestCommandHistory:
         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
+        assert history.get_next().text == expected_items[-1]
 
     def test_get_prev(self):
         commands = ["command1", "command2"]
@@ -194,7 +197,7 @@ class TestCommandHistory:
         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
+        assert history.get_prev().text == expected_items[-1]
 
 
 class TestCommandBuffer:
-- 
cgit v1.2.3