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 175d6b215..dd723c14d 100644 --- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py +++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py @@ -15,10 +15,6 @@ LOG = logging.getLogger(__name__) __author__ = 'VakarisZ' -# Default commands for executing PBA file and then removing it -DEFAULT_LINUX_COMMAND = "chmod +x {0} ; {0} ; rm {0}" -DEFAULT_WINDOWS_COMMAND = "{0} & del {0}" - DIR_CHANGE_WINDOWS = 'cd %s & ' DIR_CHANGE_LINUX = 'cd %s ; ' @@ -31,30 +27,23 @@ class UsersPBA(PBA): def __init__(self): super(UsersPBA, self).__init__(POST_BREACH_FILE_EXECUTION) self.filename = '' + if not is_windows_os(): # Add linux commands to PBA's if WormConfiguration.PBA_linux_filename: + self.filename = 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 - self.filename = WormConfiguration.PBA_linux_filename - else: - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_linux_filename) - self.command = 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 else: # Add windows commands to PBA's if WormConfiguration.PBA_windows_filename: + self.filename = 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 - self.filename = WormConfiguration.PBA_windows_filename - else: - file_path = os.path.join(get_monkey_dir_path(), WormConfiguration.PBA_windows_filename) - self.command = 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 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 000000000..83af6e00a --- /dev/null +++ b/monkey/infection_monkey/post_breach/tests/actions/test_users_custom_pba.py @@ -0,0 +1,152 @@ +import pytest + +from infection_monkey.post_breach.actions.users_custom_pba import ( + DIR_CHANGE_LINUX, DIR_CHANGE_WINDOWS, UsersPBA) + +MONKEY_DIR_PATH = "/dir/to/monkey/" +CUSTOM_LINUX_CMD = "command-for-linux" +CUSTOM_LINUX_FILENAME = "filename-for-linux" +CUSTOM_WINDOWS_CMD = "command-for-windows" +CUSTOM_WINDOWS_FILENAME = "filename-for-windows" + + +@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( + set_os_linux, fake_monkey_dir_path, monkeypatch +): + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", + CUSTOM_LINUX_CMD, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", + CUSTOM_LINUX_FILENAME, + ) + return UsersPBA() + + +def test_command_linux_custom_file_and_cmd( + mock_UsersPBA_linux_custom_file_and_cmd, +): + expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" + assert mock_UsersPBA_linux_custom_file_and_cmd.command == expected_command + + +@pytest.fixture +def mock_UsersPBA_windows_custom_file_and_cmd( + set_os_windows, fake_monkey_dir_path, monkeypatch +): + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", + CUSTOM_WINDOWS_CMD, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", + CUSTOM_WINDOWS_FILENAME, + ) + return UsersPBA() + + +def test_command_windows_custom_file_and_cmd( + mock_UsersPBA_windows_custom_file_and_cmd, +): + expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" + assert mock_UsersPBA_windows_custom_file_and_cmd.command == expected_command + + +@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_linux_custom_file(mock_UsersPBA_linux_custom_file): + expected_command = "" + assert mock_UsersPBA_linux_custom_file.command == expected_command + + +@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_windows_custom_file(mock_UsersPBA_windows_custom_file): + expected_command = "" + assert mock_UsersPBA_windows_custom_file.command == expected_command + + +@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, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_linux_filename", None + ) + return UsersPBA() + + +def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd): + expected_command = CUSTOM_LINUX_CMD + assert mock_UsersPBA_linux_custom_cmd.command == expected_command + + +@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, + ) + monkeypatch.setattr( + "infection_monkey.config.WormConfiguration.PBA_windows_filename", None + ) + return UsersPBA() + + +def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd): + expected_command = CUSTOM_WINDOWS_CMD + assert mock_UsersPBA_windows_custom_cmd.command == expected_command diff --git a/monkey/monkey_island/cc/services/config_schema/monkey.py b/monkey/monkey_island/cc/services/config_schema/monkey.py index 01d463672..82a394b65 100644 --- a/monkey/monkey_island/cc/services/config_schema/monkey.py +++ b/monkey/monkey_island/cc/services/config_schema/monkey.py @@ -11,33 +11,39 @@ MONKEY = { "type": "object", "properties": { "custom_PBA_linux_cmd": { - "title": "Linux post breach command", + "title": "Linux post-breach command", "type": "string", "default": "", - "description": "Linux command to be executed after breaching." + "description": "Command to be executed after breaching. " + "Use this field to run custom commands or execute uploaded " + "files on exploited machines.\nExample: " + "\"chmod +x ./my_script.sh; ./my_script.sh ; rm ./my_script.sh\"" }, "PBA_linux_file": { - "title": "Linux post breach file", + "title": "Linux post-breach file", "type": "string", "format": "data-url", - "description": "File to be executed after breaching. " - "If you want custom execution behavior, " - "specify it in 'Linux post breach command' field. " + "description": "File to be uploaded after breaching. " + "Use the 'Linux post-breach command' field to " + "change permissions, run, or delete the file. " "Reference your file by filename." }, "custom_PBA_windows_cmd": { - "title": "Windows post breach command", + "title": "Windows post-breach command", "type": "string", "default": "", - "description": "Windows command to be executed after breaching." + "description": "Command to be executed after breaching. " + "Use this field to run custom commands or execute uploaded " + "files on exploited machines.\nExample: " + "\"my_script.bat & del my_script.bat\"" }, "PBA_windows_file": { - "title": "Windows post breach file", + "title": "Windows post-breach file", "type": "string", "format": "data-url", - "description": "File to be executed after breaching. " - "If you want custom execution behavior, " - "specify it in 'Windows post breach command' field. " + "description": "File to be uploaded after breaching. " + "Use the 'Windows post-breach command' field to " + "change permissions, run, or delete the file. " "Reference your file by filename." }, "PBA_windows_filename": {