Fixed the logic according to some tests, added new tests

This commit is contained in:
Henrique 2019-11-24 20:13:25 -05:00
parent 1d43abcb28
commit 7b386d5393
4 changed files with 296 additions and 243 deletions

View File

@ -8,15 +8,11 @@ from mitmproxy import ctx
class CommandHistory: class CommandHistory:
def __init__(self, size: int = 300) -> None: def __init__(self, size: int = 300) -> None:
self.saved_commands: typing.Deque[str] = collections.deque( self.saved_commands: typing.Deque[str] = collections.deque(maxlen=size)
maxlen=size
)
self.index: int = 0
self.filter: str = ''
self.filtered_index: int = 0
self.filtered_commands: typing.Deque[str] = collections.deque() self.filtered_commands: typing.Deque[str] = collections.deque()
self.filter_active: bool = True self.current_index: int = -1
self.filter_str: str = ''
_command_history_path = os.path.join(os.path.expanduser(ctx.options.confdir), 'command_history') _command_history_path = os.path.join(os.path.expanduser(ctx.options.confdir), 'command_history')
_history_lines = [] _history_lines = []
@ -26,11 +22,7 @@ class CommandHistory:
self.command_history_file = open(_command_history_path, 'w') self.command_history_file = open(_command_history_path, 'w')
for l in _history_lines: for l in _history_lines:
self.add_command(l.strip(), True) self.add_command(l.strip())
@property
def last_index(self):
return len(self.saved_commands) - 1
@property @property
def last_filtered_index(self): def last_filtered_index(self):
@ -39,99 +31,63 @@ class CommandHistory:
@command.command("command_history.clear") @command.command("command_history.clear")
def clear_history(self): def clear_history(self):
self.saved_commands.clear() self.saved_commands.clear()
self.index = 0 self.filtered_commands.clear()
self.command_history_file.truncate(0) self.command_history_file.truncate(0)
self.command_history_file.seek(0) self.command_history_file.seek(0)
self.command_history_file.flush() self.command_history_file.flush()
self.filter = '' self.restart()
self.filtered_index = 0
self.filtered_commands.clear() @command.command("command_history.cancel")
self.filter_active = True def restart(self) -> None:
self.filtered_commands = self.saved_commands.copy()
self.current_index = -1
@command.command("command_history.next") @command.command("command_history.next")
def get_next(self) -> str: def get_next(self) -> str:
if self.last_index == -1:
return ''
if self.filter != '': if self.current_index == -1 or self.current_index == self.last_filtered_index:
if self.filtered_index < self.last_filtered_index: self.current_index = -1
self.filtered_index = self.filtered_index + 1 return ''
ret = self.filtered_commands[self.filtered_index] elif self.current_index < self.last_filtered_index:
else: self.current_index += 1
if self.index == -1:
ret = '' ret = self.filtered_commands[self.current_index]
elif self.index < self.last_index:
self.index = self.index + 1
ret = self.saved_commands[self.index]
else:
self.index = -1
ret = ''
return ret return ret
@command.command("command_history.prev") @command.command("command_history.prev")
def get_prev(self) -> str: def get_prev(self) -> str:
if self.last_index == -1:
if self.current_index == -1:
if self.last_filtered_index >= 0:
self.current_index = self.last_filtered_index
else:
return '' return ''
if self.filter != '': elif self.current_index > 0:
if self.filtered_index > 0: self.current_index -= 1
self.filtered_index = self.filtered_index - 1
ret = self.filtered_commands[self.filtered_index]
else:
if self.index == -1:
self.index = self.last_index
elif self.index > 0:
self.index = self.index - 1
ret = self.saved_commands[self.index] ret = self.filtered_commands[self.current_index]
return ret return ret
@command.command("command_history.filter") @command.command("command_history.filter")
def set_filter(self, command: str) -> None: def set_filter(self, command: str) -> None:
""" self.filter_str = command
This is used when the user starts typing part of a command
and then press the "up" arrow. This way, the results returned are
only for the command that the user started typing
"""
if command.strip() == '':
return
if self.filter != '':
last_filtered_command = self.filtered_commands[-1]
if command == last_filtered_command:
self.filter = ''
self.filtered_commands = []
self.filtered_index = 0
else:
self.filter = command
_filtered_commands = [c for c in self.saved_commands if c.startswith(command)] _filtered_commands = [c for c in self.saved_commands if c.startswith(command)]
self.filtered_commands = collections.deque(_filtered_commands) self.filtered_commands = collections.deque(_filtered_commands)
if command not in self.filtered_commands: if command and command not in self.filtered_commands:
self.filtered_commands.append(command) self.filtered_commands.append(command)
self.filtered_index = self.last_filtered_index self.current_index = -1
# No commands found, so act like no filter was added
if len(self.filtered_commands) == 1:
self.add_command(command)
self.filter = ''
@command.command("command_history.cancel")
def restart(self) -> None:
self.index = -1
self.filter = ''
self.filtered_commands = []
self.filtered_index = 0
@command.command("command_history.add") @command.command("command_history.add")
def add_command(self, command: str, execution: bool = False) -> None: def add_command(self, command: str) -> None:
if command.strip() == '': if command.strip() == '':
return return
if execution:
if command in self.saved_commands: if command in self.saved_commands:
self.saved_commands.remove(command) self.saved_commands.remove(command)
@ -144,6 +100,3 @@ class CommandHistory:
self.command_history_file.flush() self.command_history_file.flush()
self.restart() self.restart()
else:
if command not in self.saved_commands:
self.saved_commands.append(command)

View File

@ -153,26 +153,46 @@ class CommandEdit(urwid.WidgetWrap):
def __init__(self, master: mitmproxy.master.Master, text: str) -> None: def __init__(self, master: mitmproxy.master.Master, text: str) -> None:
super().__init__(urwid.Text(self.leader)) super().__init__(urwid.Text(self.leader))
self.master = master self.master = master
self.active_filter = False
self.filter_str = ''
self.cbuf = CommandBuffer(master, text) self.cbuf = CommandBuffer(master, text)
self.update() self.update()
def keypress(self, size, key) -> None: def keypress(self, size, key) -> None:
if key == "backspace": if key == "backspace":
self.cbuf.backspace() self.cbuf.backspace()
if self.cbuf.text == '':
self.active_filter = False
self.master.commands.execute("command_history.filter ''")
self.filter_str = ''
elif key == "left": elif key == "left":
self.cbuf.left() self.cbuf.left()
elif key == "right": elif key == "right":
self.cbuf.right() self.cbuf.right()
elif key == "up": elif key == "up":
if self.active_filter is False:
self.active_filter = True
self.filter_str = self.cbuf.text
_cmd = command_lexer.quote(self.cbuf.text) _cmd = command_lexer.quote(self.cbuf.text)
self.master.commands.execute("command_history.filter %s" % _cmd) self.master.commands.execute("command_history.filter %s" % _cmd)
cmd = self.master.commands.execute("command_history.prev") cmd = self.master.commands.execute("command_history.prev")
self.cbuf = CommandBuffer(self.master, cmd) self.cbuf = CommandBuffer(self.master, cmd)
elif key == "down": elif key == "down":
_cmd = command_lexer.quote(self.cbuf.text) prev_cmd = self.cbuf.text
self.master.commands.execute("command_history.filter %s" % _cmd)
cmd = self.master.commands.execute("command_history.next") cmd = self.master.commands.execute("command_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.filter_str = ''
self.cbuf = CommandBuffer(self.master, '')
else:
self.cbuf = CommandBuffer(self.master, cmd) self.cbuf = CommandBuffer(self.master, cmd)
elif key == "shift tab": elif key == "shift tab":
self.cbuf.cycle_completion(False) self.cbuf.cycle_completion(False)
elif key == "tab": elif key == "tab":

View File

@ -141,7 +141,7 @@ class ActionBar(urwid.WidgetWrap):
self.prompt_execute(k) self.prompt_execute(k)
elif k == "enter": elif k == "enter":
cmd = command_lexer.quote(self._w.cbuf.text) cmd = command_lexer.quote(self._w.cbuf.text)
self.master.commands.execute(f"command_history.add {cmd} true") self.master.commands.execute(f"command_history.add {cmd}")
self.prompt_execute(self._w.get_edit_text()) self.prompt_execute(self._w.get_edit_text())
else: else:
if common.is_keypress(k): if common.is_keypress(k):

View File

@ -1,9 +1,34 @@
import os
import pytest import pytest
import shutil
import uuid
from mitmproxy import options
from mitmproxy.addons import command_history
from mitmproxy.test import taddons from mitmproxy.test import taddons
from mitmproxy.tools.console.commander import commander from mitmproxy.tools.console.commander import commander
@pytest.fixture(autouse=True)
def tctx():
# 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)
opts = options.Options()
opts.set(*[f"confdir={confdir}"])
tctx = taddons.context(options=opts)
ch = command_history.CommandHistory()
tctx.master.addons.add(ch)
yield tctx
# This runs after each test
shutil.rmtree(confdir)
class TestListCompleter: class TestListCompleter:
def test_cycle(self): def test_cycle(self):
tests = [ tests = [
@ -23,53 +48,49 @@ class TestListCompleter:
["b", "ba", "bb", "b"] ["b", "ba", "bb", "b"]
], ],
] ]
for start, options, cycle in tests: for start, opts, cycle in tests:
c = commander.ListCompleter(start, options) c = commander.ListCompleter(start, opts)
for expected in cycle: for expected in cycle:
assert c.cycle() == expected assert c.cycle() == expected
class TestCommandEdit: class TestCommandEdit:
def test_open_command_bar(self):
with taddons.context() as tctx: def test_open_command_bar(self, tctx):
history = commander.CommandHistory(tctx.master, size=3) edit = commander.CommandEdit(tctx.master, '')
edit = commander.CommandEdit(tctx.master, '', history)
try: try:
edit.update() edit.update()
except IndexError: except IndexError:
pytest.faied("Unexpected IndexError") pytest.faied("Unexpected IndexError")
def test_insert(self): def test_insert(self, tctx):
with taddons.context() as tctx: edit = commander.CommandEdit(tctx.master, '')
history = commander.CommandHistory(tctx.master, size=3)
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, 'a') edit.keypress(1, 'a')
assert edit.get_edit_text() == 'a' assert edit.get_edit_text() == 'a'
# Don't let users type a space before starting a command # Don't let users type a space before starting a command
# as a usability feature # as a usability feature
history = commander.CommandHistory(tctx.master, size=3) edit = commander.CommandEdit(tctx.master, '')
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, ' ') edit.keypress(1, ' ')
assert edit.get_edit_text() == '' assert edit.get_edit_text() == ''
def test_backspace(self): def test_backspace(self, tctx):
with taddons.context() as tctx: edit = commander.CommandEdit(tctx.master, '')
history = commander.CommandHistory(tctx.master, size=3)
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, 'a') edit.keypress(1, 'a')
edit.keypress(1, 'b') edit.keypress(1, 'b')
assert edit.get_edit_text() == 'ab' assert edit.get_edit_text() == 'ab'
edit.keypress(1, 'backspace') edit.keypress(1, 'backspace')
assert edit.get_edit_text() == 'a' assert edit.get_edit_text() == 'a'
def test_left(self): def test_left(self, tctx):
with taddons.context() as tctx: edit = commander.CommandEdit(tctx.master, '')
history = commander.CommandHistory(tctx.master, size=3)
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, 'a') edit.keypress(1, 'a')
assert edit.cbuf.cursor == 1 assert edit.cbuf.cursor == 1
edit.keypress(1, 'left') edit.keypress(1, 'left')
assert edit.cbuf.cursor == 0 assert edit.cbuf.cursor == 0
@ -77,10 +98,9 @@ class TestCommandEdit:
edit.keypress(1, 'left') edit.keypress(1, 'left')
assert edit.cbuf.cursor == 0 assert edit.cbuf.cursor == 0
def test_right(self): def test_right(self, tctx):
with taddons.context() as tctx: edit = commander.CommandEdit(tctx.master, '')
history = commander.CommandHistory(tctx.master, size=3)
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, 'a') edit.keypress(1, 'a')
assert edit.cbuf.cursor == 1 assert edit.cbuf.cursor == 1
@ -91,113 +111,174 @@ class TestCommandEdit:
# Make sure cursor goes left and then back right # Make sure cursor goes left and then back right
edit.keypress(1, 'left') edit.keypress(1, 'left')
assert edit.cbuf.cursor == 0 assert edit.cbuf.cursor == 0
edit.keypress(1, 'right') edit.keypress(1, 'right')
assert edit.cbuf.cursor == 1 assert edit.cbuf.cursor == 1
def test_up_and_down(self): def test_up_and_down(self, tctx):
with taddons.context() as tctx: edit = commander.CommandEdit(tctx.master, '')
history = commander.CommandHistory(tctx.master, size=3)
history.clear_history()
edit = commander.CommandEdit(tctx.master, '', history)
buf = commander.CommandBuffer(tctx.master, 'cmd1') tctx.master.commands.execute('command_history.clear')
history.add_command(buf) tctx.master.commands.execute('command_history.add "cmd1"')
buf = commander.CommandBuffer(tctx.master, 'cmd2')
history.add_command(buf) edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit = commander.CommandEdit(tctx.master, '')
tctx.master.commands.execute('command_history.clear')
tctx.master.commands.execute('command_history.add "cmd1"')
tctx.master.commands.execute('command_history.add "cmd2"')
edit.keypress(1, 'up') edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd2' assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'up') edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1' assert edit.get_edit_text() == 'cmd1'
history = commander.CommandHistory(tctx.master, size=5) edit.keypress(1, 'up')
history.clear_history() assert edit.get_edit_text() == 'cmd1'
edit = commander.CommandEdit(tctx.master, '', history)
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'a') edit.keypress(1, 'a')
edit.keypress(1, 'b') edit.keypress(1, 'b')
edit.keypress(1, 'c') edit.keypress(1, 'c')
assert edit.get_edit_text() == 'abc' assert edit.get_edit_text() == 'abc'
edit.keypress(1, 'up') edit.keypress(1, 'up')
assert edit.get_edit_text() == ''
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'abc' assert edit.get_edit_text() == 'abc'
edit.keypress(1, 'down') edit.keypress(1, 'down')
assert edit.get_edit_text() == 'abc' assert edit.get_edit_text() == 'abc'
history = commander.CommandHistory(tctx.master, size=5) edit.keypress(1, 'down')
edit = commander.CommandEdit(tctx.master, '', history) assert edit.get_edit_text() == 'abc'
buf = commander.CommandBuffer(tctx.master, 'cmd3')
history.add_command(buf) edit.keypress(1, 'up')
assert edit.get_edit_text() == 'abc'
edit = commander.CommandEdit(tctx.master, '')
tctx.master.commands.execute('command_history.add "cmd3"')
edit.keypress(1, 'z') edit.keypress(1, 'z')
edit.keypress(1, 'up') edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd3' assert edit.get_edit_text() == 'z'
edit.keypress(1, 'down') edit.keypress(1, 'down')
assert edit.get_edit_text() == 'z' assert edit.get_edit_text() == 'z'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'z'
class TestCommandHistory: edit.keypress(1, 'backspace')
def fill_history(self, commands): assert edit.get_edit_text() == ''
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)
return history, tctx.master
def test_add_command(self): edit.keypress(1, 'up')
commands = ["command1", "command2"] assert edit.get_edit_text() == 'cmd3'
history, tctx_master = self.fill_history(commands)
saved_commands = [buf.text for buf in history.saved_commands] edit.keypress(1, 'up')
assert saved_commands == [""] + commands assert edit.get_edit_text() == 'cmd2'
# The history size is only 3. So, we forget the first edit.keypress(1, 'up')
# one command, when adding fourth command assert edit.get_edit_text() == 'cmd1'
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 edit.keypress(1, 'down')
history.add_command(cbuf) assert edit.get_edit_text() == 'cmd2'
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 edit.keypress(1, 'down')
# and replace the last command buffer if it is empty or has the same text assert edit.get_edit_text() == 'cmd3'
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): edit.keypress(1, 'down')
commands = ["command1", "command2"] assert edit.get_edit_text() == ''
history, tctx_master = self.fill_history(commands)
history.index = -1 edit.keypress(1, 'c')
expected_items = ["", "command1", "command2"] assert edit.get_edit_text() == 'c'
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().text == expected_items[-1]
def test_get_prev(self): edit.keypress(1, 'down')
commands = ["command1", "command2"] assert edit.get_edit_text() == ''
history, tctx_master = self.fill_history(commands)
expected_items = ["command2", "command1", ""] edit.keypress(1, 'up')
history.index = history.last_index + 1 assert edit.get_edit_text() == 'cmd3'
for i in range(3):
assert history.get_prev().text == expected_items[i] edit.keypress(1, 'down')
# We are at the first item of the history assert edit.get_edit_text() == ''
assert history.get_prev().text == expected_items[-1]
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
edit.keypress(1, 'backspace')
assert edit.get_edit_text() == ''
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'up')
assert edit.get_edit_text() == 'cmd1'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd2'
edit.keypress(1, 'down')
assert edit.get_edit_text() == 'cmd3'
edit.keypress(1, 'down')
assert edit.get_edit_text() == ''
class TestCommandBuffer: class TestCommandBuffer:
@ -258,8 +339,7 @@ class TestCommandBuffer:
cb.cursor = len(cb.text) cb.cursor = len(cb.text)
cb.cycle_completion() cb.cycle_completion()
ch = commander.CommandHistory(tctx.master, 30) ce = commander.CommandEdit(tctx.master, "se")
ce = commander.CommandEdit(tctx.master, "se", ch)
ce.keypress(1, 'tab') ce.keypress(1, 'tab')
ce.update() ce.update()
ret = ce.cbuf.render() ret = ce.cbuf.render()