Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

simplify command history addon

Browse Source
This commit is contained in:
Maximilian Hils 2019-11-26 02:42:47 +01:00
parent 68b016e180
commit 06ef7350f5
3 changed files with 66 additions and 131 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

177
mitmproxy/addons/command_history.py
View File

@ -1,6 +1,5 @@
import atexit
import collections
import os
import pathlib
import typing
from mitmproxy import command
@ -8,137 +7,77 @@ 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
VACUUM_SIZE = 1024
self.filtered_commands: typing.Deque[str] = collections.deque()
self.current_index: int = -1
self.filter_str: str = ''
self.command_history_path: str = ''
def __init__(self) -> None:
self.history: typing.List[str] = []
self.filtered_history: typing.List[str] = [""]
self.current_index: int = 0
atexit.register(self.cleanup)
def cleanup(self):
self._sync_saved_commands()
def load(self, loader):
loader.add_option(
"command_history", bool, True,
"""Persist command history between mitmproxy invocations."""
)
@property
def last_filtered_index(self):
return len(self.filtered_commands) - 1
def history_file(self) -> pathlib.Path:
return pathlib.Path(os.path.expanduser(ctx.options.confdir)) / "command_history"
@command.command("command_history.clear")
def clear_history(self):
self.saved_commands.clear()
self.filtered_commands.clear()
def running(self):
# FIXME: We have a weird bug where the contract for configure is not followed and it is never called with
# confdir or command_history as updated.
self.configure("command_history")
with open(self.command_history_path, 'w') as f:
f.truncate(0)
f.seek(0)
f.flush()
f.close()
def configure(self, updated):
if "command_history" in updated or "confdir" in updated:
if ctx.options.command_history and self.history_file.is_file():
self.history = self.history_file.read_text().splitlines()
self.restart()
def done(self):
if ctx.options.command_history and len(self.history) > self.VACUUM_SIZE:
# vacuum history so that it doesn't grow indefinitely.
history_str = "\n".join(self.history[-self.VACUUM_SIZE/2:]) + "\n"
self.history_file.write_text(history_str)
@command.command("command_history.cancel")
def restart(self) -> None:
self.filtered_commands = self.saved_commands.copy()
self.current_index = -1
@command.command("command_history.next")
def get_next(self) -> str:
if self.current_index == -1 or self.current_index == self.last_filtered_index:
self.current_index = -1
return ''
elif self.current_index < self.last_filtered_index:
self.current_index += 1
ret = self.filtered_commands[self.current_index]
return ret
@command.command("command_history.prev")
def get_prev(self) -> str:
if self.current_index == -1:
if self.last_filtered_index >= 0:
self.current_index = self.last_filtered_index
else:
return ''
elif self.current_index > 0:
self.current_index -= 1
ret = self.filtered_commands[self.current_index]
return ret
@command.command("command_history.filter")
def set_filter(self, command: str) -> None:
self.filter_str = command
_filtered_commands = [c for c in self.saved_commands if c.startswith(command)]
self.filtered_commands = collections.deque(_filtered_commands)
if command and command not in self.filtered_commands:
self.filtered_commands.append(command)
self.current_index = -1
@command.command("command_history.add")
@command.command("commands.history.add")
def add_command(self, command: str) -> None:
if command.strip() == '':
if not command.strip():
return
self._sync_saved_commands()
self.history.append(command)
if ctx.options.command_history:
with self.history_file.open("a") as f:
f.write(f"{command}\n")
if command in self.saved_commands:
self.saved_commands.remove(command)
self.saved_commands.append(command)
@command.command("commands.history.get")
def get_history(self) -> typing.Sequence[str]:
"""Get the entire command history."""
return self.history.copy()
_history_str = "\n".join(self.saved_commands)
with open(self.command_history_path, 'w') as f:
f.truncate(0)
f.seek(0)
f.write(_history_str)
f.flush()
f.close()
@command.command("commands.history.clear")
def clear_history(self):
self.history_file.unlink()
self.history = []
self.restart()
# Functionality to provide a filtered list that can be iterated through.
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
@command.command("commands.history.filter")
def set_filter(self, prefix: str) -> None:
self.filtered_history = [
cmd
for cmd in self.history
if cmd.startswith(prefix)
]
self.filtered_history.append(prefix)
self.current_index = len(self.filtered_history) - 1
with open(self.command_history_path, 'r') as f:
_history_lines = f.readlines()
f.close()
@command.command("commands.history.next")
def get_next(self) -> str:
self.current_index = min(self.current_index + 1, len(self.filtered_history) - 1)
return self.filtered_history[self.current_index]
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
@command.command("commands.history.prev")
def get_prev(self) -> str:
self.current_index = max(0, self.current_index - 1)
return self.filtered_history[self.current_index]

13
mitmproxy/tools/console/commander/commander.py
View File

@ -9,8 +9,6 @@ import mitmproxy.flow
import mitmproxy.master
import mitmproxy.types
from mitmproxy import command_lexer
class Completer:
@abc.abstractmethod
@ -163,7 +161,7 @@ class CommandEdit(urwid.WidgetWrap):
self.cbuf.backspace()
if self.cbuf.text == '':
self.active_filter = False
self.master.commands.execute("command_history.filter ''")
self.master.commands.call("commands.history.filter", "")
self.filter_str = ''
elif key == "left":
self.cbuf.left()
@ -173,21 +171,20 @@ class CommandEdit(urwid.WidgetWrap):
if self.active_filter is False:
self.active_filter = True
self.filter_str = self.cbuf.text
_cmd = command_lexer.quote(self.cbuf.text)
self.master.commands.execute("command_history.filter %s" % _cmd)
self.master.commands.call("commands.history.filter", self.cbuf.text)
cmd = self.master.commands.execute("command_history.prev")
cmd = self.master.commands.execute("commands.history.prev")
self.cbuf = CommandBuffer(self.master, cmd)
elif key == "down":
prev_cmd = self.cbuf.text
cmd = self.master.commands.execute("command_history.next")
cmd = self.master.commands.execute("commands.history.next")
if cmd == '':
if prev_cmd == self.filter_str:
self.cbuf = CommandBuffer(self.master, prev_cmd)
else:
self.active_filter = False
self.master.commands.execute("command_history.filter ''")
self.master.commands.call("commands.history.filter", "")
self.filter_str = ''
self.cbuf = CommandBuffer(self.master, '')
else:

7
mitmproxy/tools/console/statusbar.py
View File

@ -132,7 +132,6 @@ class ActionBar(urwid.WidgetWrap):
def keypress(self, size, k):
if self.prompting:
if k == "esc":
self.master.commands.execute('command_history.cancel')
self.prompt_done()
elif self.onekey:
if k == "enter":
@ -140,9 +139,9 @@ class ActionBar(urwid.WidgetWrap):
elif k in self.onekey:
self.prompt_execute(k)
elif k == "enter":
cmd = command_lexer.quote(self._w.cbuf.text)
self.master.commands.execute(f"command_history.add {cmd}")
self.prompt_execute(self._w.get_edit_text())
text = self._w.get_edit_text()
self.prompt_execute(text)
self.master.commands.call("commands.history.add", text)
else:
if common.is_keypress(k):
self._w.keypress(size, k)