diff --git a/src/attackmate/executors/ssh/interactfeature.py b/src/attackmate/executors/ssh/interactfeature.py index b9dbf16d5..e4aec6820 100644 --- a/src/attackmate/executors/ssh/interactfeature.py +++ b/src/attackmate/executors/ssh/interactfeature.py @@ -13,14 +13,20 @@ def __init__(self): self.logger = logging.getLogger('playbook') def check_prompt(self, output: str, prompts: list[str]) -> bool: - if output and prompts: - for p in prompts: - if output.endswith(p): - self.logger.debug('found prompt!') - self.timer = None - return True - else: + if not output: + return False + + if not prompts: self.set_timer() + return False + + for p in prompts: + if output.endswith(p): + self.logger.debug('found prompt!') + self.timer = None + return True + + self.set_timer() return False def check_timer(self, seconds: int) -> bool: diff --git a/test/units/test_interactfeature.py b/test/units/test_interactfeature.py new file mode 100644 index 000000000..40b28dd65 --- /dev/null +++ b/test/units/test_interactfeature.py @@ -0,0 +1,33 @@ +import time +from attackmate.executors.ssh.interactfeature import Interactive + + +class TestCheckPrompt: + def setup_method(self): + self.interactive = Interactive() + self.interactive.set_timer() + + def test_no_prompts_configured_resets_timer(self): + before = self.interactive.timer + time.sleep(0.05) + self.interactive.check_prompt('some output', []) + assert self.interactive.timer > before + + def test_prompt_not_matched_resets_timer(self): + before = self.interactive.timer + time.sleep(0.05) + result = self.interactive.check_prompt('still running...', ['$ ']) + assert result is False + assert self.interactive.timer > before + + def test_prompt_matched_clears_timer_and_returns_true(self): + result = self.interactive.check_prompt('root@host:~$ ', ['$ ']) + assert result is True + assert self.interactive.timer is None + + def test_empty_output_on_channel_eof_does_not_reset_timer(self): + before = self.interactive.timer + time.sleep(0.05) + result = self.interactive.check_prompt('', ['$ ']) + assert result is False + assert self.interactive.timer == before