Modify logic for user PBA command/file to allow both simultaneously

monkey/infection_monkey/post_breach/actions/users_custom_pba.py

@@ -8,6 +8,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 +32,50 @@ 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
-                    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
+            # 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
-                    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
-
-    def _execute_default(self):
+            # 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:
             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):

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