From 68b016e180ec1475391f2e8389ae1c708692c499 Mon Sep 17 00:00:00 2001 From: Henrique Date: Mon, 25 Nov 2019 14:37:49 -0500 Subject: Addressing comments from review --- mitmproxy/addons/command_history.py | 69 ++++++++++++++++++++++--------------- 1 file changed, 42 insertions(+), 27 deletions(-) (limited to 'mitmproxy') diff --git a/mitmproxy/addons/command_history.py b/mitmproxy/addons/command_history.py index fb27e805..ad4c2346 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 -- cgit v1.2.3