diff --git a/monkey/infection_monkey/exploit/powershell.py b/monkey/infection_monkey/exploit/powershell.py index 1c496fc68..8bdf7e571 100644 --- a/monkey/infection_monkey/exploit/powershell.py +++ b/monkey/infection_monkey/exploit/powershell.py @@ -1,5 +1,5 @@ import logging -import os +from pathlib import Path from typing import List, Optional from infection_monkey.exploit.HostExploiter import HostExploiter @@ -19,7 +19,7 @@ from infection_monkey.exploit.powershell_utils.powershell_client import ( IPowerShellClient, PowerShellClient, ) -from infection_monkey.exploit.tools.helpers import get_random_file_suffix +from infection_monkey.exploit.tools.helpers import get_agent_dest_path, get_random_file_suffix from infection_monkey.model import DROPPER_ARG, RUN_MONKEY, VictimHost from infection_monkey.utils.commands import build_monkey_commandline from infection_monkey.utils.environment import is_windows_os @@ -170,7 +170,7 @@ class PowerShellExploiter(HostExploiter): raise ValueError(f"Unknown secret type {credentials.secret_type}") def _execute_monkey_agent_on_victim(self): - monkey_path_on_victim = self.options["dropper_target_path_win_64"] + monkey_path_on_victim = get_agent_dest_path(self.host, self.options) self._copy_monkey_binary_to_victim(monkey_path_on_victim) logger.info("Successfully copied the monkey binary to the victim.") @@ -182,9 +182,9 @@ class PowerShellExploiter(HostExploiter): f"Failed to execute the agent binary on the victim: {ex}" ) - def _copy_monkey_binary_to_victim(self, monkey_path_on_victim): + def _copy_monkey_binary_to_victim(self, monkey_path_on_victim: Path): - temp_monkey_binary_filepath = f"monkey_temp_bin_{get_random_file_suffix()}" + temp_monkey_binary_filepath = Path(f"./monkey_temp_bin_{get_random_file_suffix()}") self._create_local_agent_file(temp_monkey_binary_filepath) @@ -194,8 +194,8 @@ class PowerShellExploiter(HostExploiter): except Exception as ex: raise RemoteAgentCopyError(f"Failed to copy the agent binary to the victim: {ex}") finally: - if os.path.isfile(temp_monkey_binary_filepath): - os.remove(temp_monkey_binary_filepath) + if temp_monkey_binary_filepath.is_file(): + temp_monkey_binary_filepath.unlink() def _create_local_agent_file(self, binary_path): agent_binary_bytes = self.agent_repository.get_agent_binary("windows") diff --git a/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py b/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py index c0ae8b260..70e82bb66 100644 --- a/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py +++ b/monkey/infection_monkey/exploit/powershell_utils/powershell_client.py @@ -1,5 +1,6 @@ import abc import logging +from pathlib import Path from typing import Optional import pypsrp @@ -63,7 +64,7 @@ class IPowerShellClient(Protocol, metaclass=abc.ABCMeta): pass @abc.abstractmethod - def copy_file(self, src: str, dest: str) -> bool: + def copy_file(self, src: Path, dest: Path) -> bool: pass @abc.abstractmethod @@ -101,9 +102,9 @@ class PowerShellClient(IPowerShellClient): output, _, _ = self._client.execute_cmd(cmd) return output - def copy_file(self, src: str, dest: str): + def copy_file(self, src: Path, dest: Path): try: - self._client.copy(src, dest) + self._client.copy(str(src), str(dest)) logger.debug(f"Successfully copied {src} to {dest} on {self._ip_addr}") except Exception as ex: logger.error(f"Failed to copy {src} to {dest} on {self._ip_addr}: {ex}") diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py index 21a0bdeb3..c88ce99d7 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -26,6 +26,9 @@ class AuthenticationErrorForTests(Exception): mock_agent_repository = MagicMock() mock_agent_repository.get_agent_binary.return_value = BytesIO(b"BINARY_EXECUTABLE") +victim_host = VictimHost("127.0.0.1") +victim_host.os["type"] = "windows" + @pytest.fixture def powershell_arguments(): @@ -39,7 +42,7 @@ def powershell_arguments(): }, } arguments = { - "host": VictimHost("127.0.0.1"), + "host": victim_host, "options": options, "current_depth": 2, "telemetry_messenger": MagicMock(), @@ -141,7 +144,7 @@ def test_successful_copy(monkeypatch, powershell_exploiter, powershell_arguments exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) - assert DROPPER_TARGET_PATH_64 in mock_client.return_value.copy_file.call_args[0][1] + assert DROPPER_TARGET_PATH_64 in str(mock_client.return_value.copy_file.call_args[0][1]) assert exploit_result.exploitation_success