diff --git a/mitmproxy/addons/command_history.py b/mitmproxy/addons/command_history.py index fb27e8051..ad4c23468 100644 --- a/mitmproxy/addons/command_history.py +++ b/mitmproxy/addons/command_history.py @@ -10,31 +10,17 @@ from mitmproxy import ctx class CommandHistory: def __init__(self, size: int = 300) -> None: self.saved_commands: typing.Deque[str] = collections.deque(maxlen=size) + self.is_configured = False self.filtered_commands: typing.Deque[str] = collections.deque() self.current_index: int = -1 self.filter_str: str = '' - - _command_history_dir = os.path.expanduser(ctx.options.confdir) - if not os.path.exists(_command_history_dir): - os.makedirs(_command_history_dir) - - self.command_history_path = os.path.join(_command_history_dir, 'command_history') - _history_lines: typing.List[str] = [] - if os.path.exists(self.command_history_path): - _history_lines = open(self.command_history_path, 'r').readlines() - - self.command_history_file = open(self.command_history_path, 'w') - - for l in _history_lines: - self.add_command(l.strip()) + self.command_history_path: str = '' atexit.register(self.cleanup) def cleanup(self): - self._reload_saved_commands() - if self.command_history_file and not self.command_history_file.closed: - self.command_history_file.close() + self._sync_saved_commands() @property def last_filtered_index(self): @@ -44,9 +30,13 @@ class CommandHistory: def clear_history(self): self.saved_commands.clear() self.filtered_commands.clear() - self.command_history_file.truncate(0) - self.command_history_file.seek(0) - self.command_history_file.flush() + + with open(self.command_history_path, 'w') as f: + f.truncate(0) + f.seek(0) + f.flush() + f.close() + self.restart() @command.command("command_history.cancel") @@ -100,30 +90,55 @@ class CommandHistory: if command.strip() == '': return - self._reload_saved_commands() + self._sync_saved_commands() if command in self.saved_commands: self.saved_commands.remove(command) self.saved_commands.append(command) _history_str = "\n".join(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() + with open(self.command_history_path, 'w') as f: + f.truncate(0) + f.seek(0) + f.write(_history_str) + f.flush() + f.close() self.restart() - def _reload_saved_commands(self): + def _sync_saved_commands(self): # First read all commands from the file to merge anything that may # have come from a different instance of the mitmproxy or sister tools if not os.path.exists(self.command_history_path): return - _history_lines = open(self.command_history_path, 'r').readlines() + with open(self.command_history_path, 'r') as f: + _history_lines = f.readlines() + f.close() + self.saved_commands.clear() for l in _history_lines: l = l.strip() if l in self.saved_commands: self.saved_commands.remove(l) self.saved_commands.append(l.strip()) + + def configure(self, updated: typing.Set[str]): + if self.is_configured: + return + + _command_history_dir = os.path.expanduser(ctx.options.confdir) + if not os.path.exists(_command_history_dir): + os.makedirs(_command_history_dir) + + self.command_history_path = os.path.join(_command_history_dir, 'command_history') + _history_lines: typing.List[str] = [] + if os.path.exists(self.command_history_path): + with open(self.command_history_path, 'r') as f: + _history_lines = f.readlines() + f.close() + + for l in _history_lines: + self.add_command(l.strip()) + + self.is_configured = True diff --git a/test/mitmproxy/addons/test_command_history.py b/test/mitmproxy/addons/test_command_history.py index d22e1fb19..026ce53e2 100644 --- a/test/mitmproxy/addons/test_command_history.py +++ b/test/mitmproxy/addons/test_command_history.py @@ -1,7 +1,5 @@ import os import pytest -import shutil -import uuid from mitmproxy import options from mitmproxy.addons import command_history @@ -9,24 +7,22 @@ from mitmproxy.test import taddons @pytest.fixture(autouse=True) -def tctx(): +def tctx(tmpdir): # This runs before each test - dir_id = str(uuid.uuid4()) - confdir = os.path.expanduser(f"~/.mitmproxy-test-suite-{dir_id}") - if not os.path.exists(confdir): - os.makedirs(confdir) + dir_name = tmpdir.mkdir('mitmproxy').dirname + confdir = dir_name opts = options.Options() opts.set(*[f"confdir={confdir}"]) tctx = taddons.context(options=opts) ch = command_history.CommandHistory() tctx.master.addons.add(ch) + ch.configure([]) yield tctx # This runs after each test ch.cleanup() - shutil.rmtree(confdir) class TestCommandHistory: @@ -38,6 +34,7 @@ class TestCommandHistory: f.close() history = command_history.CommandHistory() + history.configure([]) saved_commands = [cmd for cmd in history.saved_commands] assert saved_commands == ['cmd1', 'cmd2', 'cmd3'] @@ -46,6 +43,7 @@ class TestCommandHistory: def test_add_command(self, tctx): history = command_history.CommandHistory(3) + history.configure([]) history.add_command('cmd1') history.add_command('cmd2') @@ -81,6 +79,7 @@ class TestCommandHistory: def test_get_next_and_prev(self, tctx): history = command_history.CommandHistory(5) + history.configure([]) history.add_command('cmd1') @@ -156,6 +155,7 @@ class TestCommandHistory: def test_clear(self, tctx): history = command_history.CommandHistory(3) + history.configure([]) history.add_command('cmd1') history.add_command('cmd2') @@ -173,6 +173,7 @@ class TestCommandHistory: def test_filter(self, tctx): history = command_history.CommandHistory(3) + history.configure([]) history.add_command('cmd1') history.add_command('cmd2') @@ -206,7 +207,6 @@ class TestCommandHistory: history.cleanup() def test_multiple_instances(self, tctx): - instances = [ command_history.CommandHistory(10), command_history.CommandHistory(10), @@ -214,6 +214,7 @@ class TestCommandHistory: ] for i in instances: + i.configure([]) saved_commands = [cmd for cmd in i.saved_commands] assert saved_commands == [] @@ -253,6 +254,7 @@ class TestCommandHistory: assert saved_commands == ['cmd1', 'cmd2', 'cmd3', 'cmd4'] instances.append(command_history.CommandHistory(10)) + instances[3].configure([]) saved_commands = [cmd for cmd in instances[3].saved_commands] assert saved_commands == ['cmd1', 'cmd2', 'cmd3', 'cmd4'] @@ -274,3 +276,31 @@ class TestCommandHistory: lines = open(_path, 'r').readlines() saved_commands = [cmd.strip() for cmd in lines] assert saved_commands == ['cmd1', 'cmd2', 'cmd3', 'cmd4', 'cmd_before_close', 'new_cmd'] + + instances = [ + command_history.CommandHistory(10), + command_history.CommandHistory(10) + ] + + for i in instances: + i.configure([]) + i.clear_history() + saved_commands = [cmd for cmd in i.saved_commands] + assert saved_commands == [] + + instances[0].add_command('cmd1') + instances[0].add_command('cmd2') + instances[1].add_command('cmd3') + instances[1].add_command('cmd4') + instances[1].add_command('cmd5') + + saved_commands = [cmd for cmd in instances[1].saved_commands] + assert saved_commands == ['cmd1', 'cmd2', 'cmd3', 'cmd4', 'cmd5'] + + instances.pop() + instances.pop() + + _path = os.path.join(tctx.options.confdir, 'command_history') + lines = open(_path, 'r').readlines() + saved_commands = [cmd.strip() for cmd in lines] + assert saved_commands == ['cmd1', 'cmd2', 'cmd3', 'cmd4', 'cmd5'] diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py index 06e676c0b..4fa10eb80 100644 --- a/test/mitmproxy/tools/console/test_commander.py +++ b/test/mitmproxy/tools/console/test_commander.py @@ -1,7 +1,4 @@ -import os import pytest -import shutil -import uuid from mitmproxy import options from mitmproxy.addons import command_history @@ -10,24 +7,22 @@ from mitmproxy.tools.console.commander import commander @pytest.fixture(autouse=True) -def tctx(): +def tctx(tmpdir): # This runs before each test - dir_id = str(uuid.uuid4()) - confdir = os.path.expanduser(f"~/.mitmproxy-test-suite-{dir_id}") - if not os.path.exists(confdir): - os.makedirs(confdir) + dir_name = tmpdir.mkdir('mitmproxy').dirname + confdir = dir_name opts = options.Options() opts.set(*[f"confdir={confdir}"]) tctx = taddons.context(options=opts) ch = command_history.CommandHistory() tctx.master.addons.add(ch) + ch.configure([]) yield tctx # This runs after each test - ch.command_history_file.close() - shutil.rmtree(confdir) + ch.cleanup() class TestListCompleter: