Merge pull request #1027 from shreyamalviya/rephrasing-config-custom-pba

Rephrase custom PBA file descriptions in configuration

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
 

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

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. "
+                    "description": "File to be uploaded after breaching. "
-                                   "If you want custom execution behavior, "
+                                   "Use the 'Linux post-breach command' field to "
-                                   "specify it in 'Linux post breach command' field. "
+                                   "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. "
+                    "description": "File to be uploaded after breaching. "
-                                   "If you want custom execution behavior, "
+                                   "Use the 'Windows post-breach command' field to "
-                                   "specify it in 'Windows post breach command' field. "
+                                   "change permissions, run, or delete the file. "
                                    "Reference your file by filename."
                 },
                 "PBA_windows_filename": {