From fc3c909ad74c35b08d81fb093ed289b703f35f21 Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 5 Mar 2021 22:58:36 +0530 Subject: [PATCH 1/4] Modify logic for user PBA command/file to allow both simultaneously --- .../post_breach/actions/users_custom_pba.py | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py index 175d6b215d9..3bb001539bf 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py @@ -1,5 +1,6 @@ import logging import os +import subprocess from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION from common.utils.attack_utils import ScanStatus @@ -8,6 +9,7 @@ from infection_monkey.network.tools import get_interface_to_target from infection_monkey.post_breach.pba import PBA from infection_monkey.telemetry.attack.t1105_telem import T1105Telem +from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.monkey_dir import get_monkey_dir_path @@ -31,37 +33,49 @@ class UsersPBA(PBA): def __init__(self): super(UsersPBA, self).__init__(POST_BREACH_FILE_EXECUTION) self.filename = '' + self.command_list = [] if not is_windows_os(): # Add linux commands to PBA's if WormConfiguration.PBA_linux_filename: if WormConfiguration.custom_PBA_linux_cmd: # Add change dir command, because user will try to access his file - self.command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd + cmd = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd + self.command_list.append(cmd) self.filename = WormConfiguration.PBA_linux_filename + if self.filename not in cmd: # PBA command is not about uploaded PBA file + file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename) + self.command_list.append(DEFAULT_LINUX_COMMAND.format(file_path)) else: file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename) - self.command = DEFAULT_LINUX_COMMAND.format(file_path) + self.command_list.append(DEFAULT_LINUX_COMMAND.format(file_path)) self.filename = WormConfiguration.PBA_linux_filename elif WormConfiguration.custom_PBA_linux_cmd: - self.command = WormConfiguration.custom_PBA_linux_cmd + self.command_list.append(WormConfiguration.custom_PBA_linux_cmd) else: # Add windows commands to PBA's if WormConfiguration.PBA_windows_filename: if WormConfiguration.custom_PBA_windows_cmd: # Add change dir command, because user will try to access his file - self.command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd + cmd = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd + self.command_list.append(cmd) self.filename = WormConfiguration.PBA_windows_filename + if self.filename not in cmd: # PBA command is not about uploaded PBA file + file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename) + self.command_list.append(DEFAULT_WINDOWS_COMMAND.format(file_path)) else: file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename) - self.command = DEFAULT_WINDOWS_COMMAND.format(file_path) + self.command_list.append(DEFAULT_WINDOWS_COMMAND.format(file_path)) self.filename = WormConfiguration.PBA_windows_filename elif WormConfiguration.custom_PBA_windows_cmd: - self.command = WormConfiguration.custom_PBA_windows_cmd + self.command_list.append(WormConfiguration.custom_PBA_windows_cmd) - def _execute_default(self): + def run(self): if self.filename: UsersPBA.download_pba_file(get_monkey_dir_path(), self.filename) - return super(UsersPBA, self)._execute_default() + for command in self.command_list: + self.command = command # needs to be an object variable since PostBreachTelem needs it + result = self._execute_default() + PostBreachTelem(self, result).send() @staticmethod def should_run(class_name): From 8d9cccc48d8a4f9e94dba157e1fcd0f0e9f74782 Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 5 Mar 2021 23:01:43 +0530 Subject: [PATCH 2/4] Remove unused import --- monkey/infection_monkey/post_breach/actions/users_custom_pba.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py index 3bb001539bf..4d7504f7628 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py @@ -1,6 +1,5 @@ import logging import os -import subprocess from common.common_consts.post_breach_consts import POST_BREACH_FILE_EXECUTION from common.utils.attack_utils import ScanStatus From d96b5e624a0f11ffbebf55c27ac617e48af64355 Mon Sep 17 00:00:00 2001 From: Shreya Date: Mon, 8 Mar 2021 13:52:40 +0530 Subject: [PATCH 3/4] Refactor setting commands in UsersPBA --- .../post_breach/actions/users_custom_pba.py | 65 ++++++++++--------- 1 file changed, 33 insertions(+), 32 deletions(-) diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py index 4d7504f7628..a5e75df2e77 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py @@ -34,39 +34,40 @@ def __init__(self): self.filename = '' self.command_list = [] if not is_windows_os(): - # Add linux commands to PBA's - if WormConfiguration.PBA_linux_filename: - if WormConfiguration.custom_PBA_linux_cmd: - # Add change dir command, because user will try to access his file - cmd = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + WormConfiguration.custom_PBA_linux_cmd - self.command_list.append(cmd) - self.filename = WormConfiguration.PBA_linux_filename - if self.filename not in cmd: # PBA command is not about uploaded PBA file - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename) - self.command_list.append(DEFAULT_LINUX_COMMAND.format(file_path)) - else: - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename) - self.command_list.append(DEFAULT_LINUX_COMMAND.format(file_path)) - self.filename = WormConfiguration.PBA_linux_filename - elif WormConfiguration.custom_PBA_linux_cmd: - self.command_list.append(WormConfiguration.custom_PBA_linux_cmd) + # Add linux commands to PBAs + self._set_default_commands(is_windows=False) + self._set_command_list(custom_filename=WormConfiguration.PBA_linux_filename, + custom_cmd=WormConfiguration.custom_PBA_linux_cmd) else: - # Add windows commands to PBA's - if WormConfiguration.PBA_windows_filename: - if WormConfiguration.custom_PBA_windows_cmd: - # Add change dir command, because user will try to access his file - cmd = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + WormConfiguration.custom_PBA_windows_cmd - self.command_list.append(cmd) - self.filename = WormConfiguration.PBA_windows_filename - if self.filename not in cmd: # PBA command is not about uploaded PBA file - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename) - self.command_list.append(DEFAULT_WINDOWS_COMMAND.format(file_path)) - else: - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename) - self.command_list.append(DEFAULT_WINDOWS_COMMAND.format(file_path)) - self.filename = WormConfiguration.PBA_windows_filename - elif WormConfiguration.custom_PBA_windows_cmd: - self.command_list.append(WormConfiguration.custom_PBA_windows_cmd) + # Add windows commands to PBAs + self._set_default_commands(is_windows=True) + self._set_command_list(custom_filename=WormConfiguration.PBA_windows_filename, + custom_cmd=WormConfiguration.custom_PBA_windows_cmd) + + def _set_default_commands(self, is_windows): + if is_windows: + self.dir_change_command = DIR_CHANGE_WINDOWS + self.default_command = DEFAULT_WINDOWS_COMMAND + else: + self.dir_change_command = DIR_CHANGE_LINUX + self.default_command = DEFAULT_LINUX_COMMAND + + def _set_command_list(self, custom_filename, custom_cmd): + if custom_filename: + if custom_cmd: + # Add change dir command, because user will try to access his file + cmd = (self.dir_change_command % get_monkey_dir_path()) + custom_cmd + self.command_list.append(cmd) + self.filename = custom_filename + if self.filename not in cmd: # PBA command is not about uploaded PBA file + file_path = os.path.join(get_monkey_dir_path(), self.filename) + self.command_list.append(self.default_command.format(file_path)) + else: + file_path = os.path.join(get_monkey_dir_path(), custom_filename) + self.command_list.append(self.default_command.format(file_path)) + self.filename = custom_filename + elif custom_cmd: + self.command_list.append(custom_cmd) def run(self): if self.filename: From 34f6b052f688ce9c3b709b7e47c05a1a1dc011d4 Mon Sep 17 00:00:00 2001 From: Shreya Date: Mon, 8 Mar 2021 13:53:01 +0530 Subject: [PATCH 4/4] Add unit tests --- .../tests/actions/test_users_custom_pba.py | 227 ++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 monkey/infection_monkey/post_breach/tests/actions/test_users_custom_pba.py diff --git a/monkey/infection_monkey/post_breach/tests/actions/test_users_custom_pba.py b/monkey/infection_monkey/post_breach/tests/actions/test_users_custom_pba.py new file mode 100644 index 00000000000..dab6ee59ef2 --- /dev/null +++ b/monkey/infection_monkey/post_breach/tests/actions/test_users_custom_pba.py @@ -0,0 +1,227 @@ +import pytest + +from infection_monkey.post_breach.actions.users_custom_pba import ( + DEFAULT_LINUX_COMMAND, DEFAULT_WINDOWS_COMMAND, DIR_CHANGE_LINUX, + DIR_CHANGE_WINDOWS, UsersPBA) + +MONKEY_DIR_PATH = "/dir/to/monkey/" +CUSTOM_LINUX_CMD_SEPARATE = "command-for-linux" +CUSTOM_LINUX_FILENAME = "filename-for-linux" +CUSTOM_LINUX_CMD_RELATED = f"command-with-{CUSTOM_LINUX_FILENAME}" +CUSTOM_WINDOWS_CMD_SEPARATE = "command-for-windows" +CUSTOM_WINDOWS_FILENAME = "filename-for-windows" +CUSTOM_WINDOWS_CMD_RELATED = f"command-with-{CUSTOM_WINDOWS_FILENAME}" + + +@pytest.fixture +def fake_monkey_dir_path(monkeypatch): + monkeypatch.setattr( + "infection_monkey.post_breach.actions.users_custom_pba.get_monkey_dir_path", + lambda: MONKEY_DIR_PATH, + ) + + +@pytest.fixture +def set_os_linux(monkeypatch): + monkeypatch.setattr( + "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + lambda: False, + ) + + +@pytest.fixture +def set_os_windows(monkeypatch): + monkeypatch.setattr( + "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + lambda: True, + ) + + +@pytest.fixture +def mock_UsersPBA_linux_custom_file_and_cmd_separate( + set_os_linux, fake_monkey_dir_path, monkeypatch +): + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", + CUSTOM_LINUX_CMD_SEPARATE, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", + CUSTOM_LINUX_FILENAME, + ) + return UsersPBA() + + +def test_command_list_linux_custom_file_and_cmd_separate( + mock_UsersPBA_linux_custom_file_and_cmd_separate, +): + expected_command_list = [ + f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD_SEPARATE}", + f"chmod +x {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME} ; {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME} ; " + + f"rm {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME}", + ] + assert ( + mock_UsersPBA_linux_custom_file_and_cmd_separate.command_list + == expected_command_list + ) + + +@pytest.fixture +def mock_UsersPBA_windows_custom_file_and_cmd_separate( + set_os_windows, fake_monkey_dir_path, monkeypatch +): + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", + CUSTOM_WINDOWS_CMD_SEPARATE, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", + CUSTOM_WINDOWS_FILENAME, + ) + return UsersPBA() + + +def test_command_list_windows_custom_file_and_cmd_separate( + mock_UsersPBA_windows_custom_file_and_cmd_separate, +): + expected_command_list = [ + f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD_SEPARATE}", + f"{MONKEY_DIR_PATH}{CUSTOM_WINDOWS_FILENAME} & del {MONKEY_DIR_PATH}{CUSTOM_WINDOWS_FILENAME}", + ] + assert ( + mock_UsersPBA_windows_custom_file_and_cmd_separate.command_list + == expected_command_list + ) + + +@pytest.fixture +def mock_UsersPBA_linux_custom_file_and_cmd_related( + set_os_linux, fake_monkey_dir_path, monkeypatch +): + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", + CUSTOM_LINUX_CMD_RELATED, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", + CUSTOM_LINUX_FILENAME, + ) + return UsersPBA() + + +def test_command_list_linux_custom_file_and_cmd_related( + mock_UsersPBA_linux_custom_file_and_cmd_related, +): + expected_command_list = [f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD_RELATED}"] + assert ( + mock_UsersPBA_linux_custom_file_and_cmd_related.command_list + == expected_command_list + ) + + +@pytest.fixture +def mock_UsersPBA_windows_custom_file_and_cmd_related( + set_os_windows, fake_monkey_dir_path, monkeypatch +): + + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", + CUSTOM_WINDOWS_CMD_RELATED, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", + CUSTOM_WINDOWS_FILENAME, + ) + return UsersPBA() + + +def test_command_list_windows_custom_file_and_cmd_related( + mock_UsersPBA_windows_custom_file_and_cmd_related, +): + expected_command_list = [ + f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD_RELATED}", + ] + assert ( + mock_UsersPBA_windows_custom_file_and_cmd_related.command_list + == expected_command_list + ) + + +@pytest.fixture +def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypatch): + + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", + CUSTOM_LINUX_FILENAME, + ) + return UsersPBA() + + +def test_command_list_linux_custom_file(mock_UsersPBA_linux_custom_file): + expected_command_list = [ + f"chmod +x {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME} ; {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME} ; " + + f"rm {MONKEY_DIR_PATH}{CUSTOM_LINUX_FILENAME}" + ] + + assert mock_UsersPBA_linux_custom_file.command_list == expected_command_list + + +@pytest.fixture +def mock_UsersPBA_windows_custom_file( + set_os_windows, fake_monkey_dir_path, monkeypatch +): + + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", + CUSTOM_WINDOWS_FILENAME, + ) + return UsersPBA() + + +def test_command_list_windows_custom_file(mock_UsersPBA_windows_custom_file): + expected_command_list = [ + f"{MONKEY_DIR_PATH}{CUSTOM_WINDOWS_FILENAME} & del {MONKEY_DIR_PATH}{CUSTOM_WINDOWS_FILENAME}", + ] + assert mock_UsersPBA_windows_custom_file.command_list == expected_command_list + + +@pytest.fixture +def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch): + + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", + CUSTOM_LINUX_CMD_SEPARATE, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", None + ) + return UsersPBA() + + +def test_command_list_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd): + expected_command_list = [CUSTOM_LINUX_CMD_SEPARATE] + assert mock_UsersPBA_linux_custom_cmd.command_list == expected_command_list + + +@pytest.fixture +def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch): + + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", + CUSTOM_WINDOWS_CMD_SEPARATE, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", None + ) + return UsersPBA() + + +def test_command_list_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd): + expected_command_list = [CUSTOM_WINDOWS_CMD_SEPARATE] + assert mock_UsersPBA_windows_custom_cmd.command_list == expected_command_list