From a12d964e1788c2bdf3b9bc96e3f616acc29fe4e4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 11:12:58 -0400 Subject: [PATCH 01/10] UT: Add InMemoryAgentBinaryRepository --- monkey/tests/monkey_island/__init__.py | 1 + .../in_memory_agent_binary_repository.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 monkey/tests/monkey_island/in_memory_agent_binary_repository.py diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index d7e2d24b0dc..e9ae5fd4441 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -12,3 +12,4 @@ from .in_memory_agent_event_repository import InMemoryAgentEventRepository from .in_memory_agent_plugin_repository import InMemoryAgentPluginRepository from .in_memory_agent_configuration_service import InMemoryAgentConfigurationService +from .in_memory_agent_binary_repository import InMemoryAgentBinaryRepository diff --git a/monkey/tests/monkey_island/in_memory_agent_binary_repository.py b/monkey/tests/monkey_island/in_memory_agent_binary_repository.py new file mode 100644 index 00000000000..00232903c65 --- /dev/null +++ b/monkey/tests/monkey_island/in_memory_agent_binary_repository.py @@ -0,0 +1,26 @@ +import io + +from common import OperatingSystem +from monkey_island.cc.services.agent_binary_service.i_agent_binary_repository import ( + IAgentBinaryRepository, +) + +LINUX_AGENT_BINARY = b"linux_binary" +WINDOWS_AGENT_BINARY = b"windows_binary" + + +class InMemoryAgentBinaryRepository(IAgentBinaryRepository): + def __init__(self): + self.agent_binaries = { + OperatingSystem.LINUX: LINUX_AGENT_BINARY, + OperatingSystem.WINDOWS: WINDOWS_AGENT_BINARY, + } + + def get_agent_binary(self, operating_system: OperatingSystem): + return io.BytesIO(self.agent_binaries[operating_system]) + + def get_linux_binary(self): + return io.BytesIO(self.agent_binaries[OperatingSystem.LINUX]) + + def get_windows_binary(self): + return io.BytesIO(self.agent_binaries[OperatingSystem.WINDOWS]) From fb515e05b32d7cd02dadbca98ce187c28ef25421 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 10:51:55 -0400 Subject: [PATCH 02/10] Island: Add IAgentBinaryRepository.get_agent_binary() --- .../agent_binary_repository.py | 9 +++++ .../i_agent_binary_repository.py | 12 +++++++ ...erade_agent_binary_repository_decorator.py | 7 ++++ .../test_agent_binary_repository.py | 35 ++++++++++++++++--- ...erade_agent_binary_repository_decorator.py | 26 +++++++------- 5 files changed, 70 insertions(+), 19 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py index 8b5d6365a68..084e1fe101b 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py @@ -1,5 +1,6 @@ from typing import BinaryIO +from common import OperatingSystem from monkey_island.cc.repositories import IFileRepository, RetrievalError from .i_agent_binary_repository import IAgentBinaryRepository @@ -7,11 +8,19 @@ LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" +AGENT_FILE_NAMES = { + OperatingSystem.LINUX: LINUX_AGENT_FILE_NAME, + OperatingSystem.WINDOWS: WINDOWS_AGENT_FILE_NAME, +} + class AgentBinaryRepository(IAgentBinaryRepository): def __init__(self, file_repository: IFileRepository): self._file_repository = file_repository + def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + return self._get_binary(AGENT_FILE_NAMES[operating_system]) + def get_linux_binary(self) -> BinaryIO: return self._get_binary(LINUX_AGENT_FILE_NAME) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py index 6c9abcc40d9..a75fbe34310 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py @@ -1,12 +1,24 @@ import abc from typing import BinaryIO +from common import OperatingSystem + class IAgentBinaryRepository(metaclass=abc.ABCMeta): """ A repository that retrieves the agent binaries """ + @abc.abstractmethod + def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + """ + Retrieve an agent binary for the specified operating system + + :param operating_system: The operating system with which the binary must be compatible + :return: A file-like object that represents the agent binary + :raises RetrievalError: If the agent binary could not be retrieved + """ + @abc.abstractmethod def get_linux_binary(self) -> BinaryIO: """ diff --git a/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py b/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py index cb429ad4400..43142aaf8eb 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py @@ -2,6 +2,8 @@ from functools import lru_cache from typing import BinaryIO +from common import OperatingSystem + from .i_agent_binary_repository import IAgentBinaryRepository DEFAULT_NULL_BYTES_LENGTH = 16 @@ -27,6 +29,11 @@ def __init__( self._masque = masque self._null_bytes_length = null_bytes_length + @lru_cache() + def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + agent_binary = self._agent_binary_repository.get_agent_binary(operating_system) + return self._apply_masque(agent_binary) + @lru_cache() def get_linux_binary(self) -> BinaryIO: agent_linux_binary = self._agent_binary_repository.get_linux_binary() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py index dac60f98541..43f812ab95e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py @@ -1,9 +1,14 @@ +import io from unittest.mock import MagicMock import pytest +from tests.monkey_island import InMemoryFileRepository +from common import OperatingSystem from monkey_island.cc.repositories import IFileRepository, RetrievalError from monkey_island.cc.services.agent_binary_service.agent_binary_repository import ( + LINUX_AGENT_FILE_NAME, + WINDOWS_AGENT_FILE_NAME, AgentBinaryRepository, ) @@ -11,6 +16,26 @@ WINDOWS_AGENT_BINARY = b"windows_binary" +@pytest.fixture +def agent_binary_repository(): + file_repository = InMemoryFileRepository() + file_repository.save_file(LINUX_AGENT_FILE_NAME, io.BytesIO(LINUX_AGENT_BINARY)) + file_repository.save_file(WINDOWS_AGENT_FILE_NAME, io.BytesIO(WINDOWS_AGENT_BINARY)) + return AgentBinaryRepository(file_repository) + + +@pytest.mark.parametrize( + "operating_system,expected_agent_binary", + ([OperatingSystem.LINUX, LINUX_AGENT_BINARY], [OperatingSystem.WINDOWS, WINDOWS_AGENT_BINARY]), +) +def test_get_os_binary( + agent_binary_repository, operating_system: OperatingSystem, expected_agent_binary: bytes +): + assert ( + agent_binary_repository.get_agent_binary(operating_system).read() == expected_agent_binary + ) + + @pytest.fixture def error_mock_file_repository(): mock_file_repository = MagicMock(wraps=IFileRepository) @@ -20,15 +45,15 @@ def error_mock_file_repository(): @pytest.fixture -def agent_binary_repository(error_mock_file_repository): +def error_raising_agent_binary_repository(error_mock_file_repository): return AgentBinaryRepository(error_mock_file_repository) -def test_get_linux_binary_retrieval_error(agent_binary_repository): +def test_get_linux_binary_retrieval_error(error_raising_agent_binary_repository): with pytest.raises(RetrievalError): - agent_binary_repository.get_linux_binary() + error_raising_agent_binary_repository.get_linux_binary() -def test_get_windows_binary_retrieval_error(agent_binary_repository): +def test_get_windows_binary_retrieval_error(error_raising_agent_binary_repository): with pytest.raises(RetrievalError): - agent_binary_repository.get_windows_binary() + error_raising_agent_binary_repository.get_windows_binary() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py index 0d0a93f1a50..8eb636a2730 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py @@ -1,8 +1,9 @@ import io -from unittest.mock import MagicMock import pytest +from tests.monkey_island import InMemoryAgentBinaryRepository +from common import OperatingSystem from monkey_island.cc.services.agent_binary_service.i_agent_binary_repository import ( IAgentBinaryRepository, ) @@ -27,21 +28,16 @@ @pytest.fixture -def mock_agent_binary_repository() -> IAgentBinaryRepository: - mock_agent_binary_repository = MagicMock(wraps=IAgentBinaryRepository) - - mock_agent_binary_repository.get_linux_binary = MagicMock(return_value=LINUX_AGENT_BINARY) - mock_agent_binary_repository.get_windows_binary = MagicMock(return_value=WINDOWS_AGENT_BINARY) - - return mock_agent_binary_repository +def in_memory_agent_binary_repository() -> InMemoryAgentBinaryRepository: + return InMemoryAgentBinaryRepository() @pytest.fixture def mock_masquerade_agent_binary_repository( - mock_agent_binary_repository, + in_memory_agent_binary_repository: IAgentBinaryRepository, ) -> MasqueradeAgentBinaryRepositoryDecorator: return MasqueradeAgentBinaryRepositoryDecorator( - mock_agent_binary_repository, MASQUE, NULL_BYTES_LENGTH + in_memory_agent_binary_repository, MASQUE, NULL_BYTES_LENGTH ) @@ -62,22 +58,24 @@ def test_masquerade_agent_binary_repository_decorator__windows( def test_masquerade_agent_binary_repository_decorator__cached_linux( - mock_agent_binary_repository: IAgentBinaryRepository, + in_memory_agent_binary_repository: InMemoryAgentBinaryRepository, mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, ): actual_linux_binary = mock_masquerade_agent_binary_repository.get_linux_binary() - mock_agent_binary_repository.get_linux_binary = MagicMock(return_value=b"not_linux_binary") # type: ignore[assignment,method-assign] # noqa: E501 + in_memory_agent_binary_repository.agent_binaries[OperatingSystem.LINUX] = b"not_linux_binary" cached_linux_binary = mock_masquerade_agent_binary_repository.get_linux_binary() assert actual_linux_binary == cached_linux_binary def test_masquerade_agent_binary_repository_decorator__cached_windows( - mock_agent_binary_repository: IAgentBinaryRepository, + in_memory_agent_binary_repository: InMemoryAgentBinaryRepository, mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, ): actual_windows_binary = mock_masquerade_agent_binary_repository.get_windows_binary() - mock_agent_binary_repository.get_windows_binary = MagicMock(return_value=b"not_windows_binary") # type: ignore[assignment,method-assign] # noqa: E501 + in_memory_agent_binary_repository.agent_binaries[ + OperatingSystem.WINDOWS + ] = b"not_windows_binary" cached_windows_binary = mock_masquerade_agent_binary_repository.get_windows_binary() assert actual_windows_binary == cached_windows_binary From c647f2ac2727880fc3c15dea1ab4507d301fb314 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 11:23:58 -0400 Subject: [PATCH 03/10] UT: Parametrize MasqueradeAgentBinaryRepositoryDecorator tests --- ...erade_agent_binary_repository_decorator.py | 49 ++++++++----------- 1 file changed, 21 insertions(+), 28 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py index 8eb636a2730..67484c689d4 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py @@ -41,41 +41,34 @@ def mock_masquerade_agent_binary_repository( ) -def test_masquerade_agent_binary_repository_decorator__linux( - mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, -): - actual_linux_binary = mock_masquerade_agent_binary_repository.get_linux_binary() - - assert actual_linux_binary.getvalue() == MASQUED_LINUX_AGENT_BINARY.getvalue() # type: ignore[attr-defined] # noqa: E501 - - -def test_masquerade_agent_binary_repository_decorator__windows( +@pytest.mark.parametrize( + "operating_system,expected_agent_binary", + ( + (OperatingSystem.LINUX, MASQUED_LINUX_AGENT_BINARY), + (OperatingSystem.WINDOWS, MASQUED_WINDOWS_AGENT_BINARY), + ), +) +def test_get_agent_binary( mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, + operating_system: OperatingSystem, + expected_agent_binary: io.BytesIO, ): - actual_windows_binary = mock_masquerade_agent_binary_repository.get_windows_binary() + actual_linux_binary = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) - assert actual_windows_binary.getvalue() == MASQUED_WINDOWS_AGENT_BINARY.getvalue() # type: ignore[attr-defined] # noqa: E501 + assert actual_linux_binary.getvalue() == expected_agent_binary.getvalue() # type: ignore[attr-defined] # noqa: E501 -def test_masquerade_agent_binary_repository_decorator__cached_linux( +@pytest.mark.parametrize( + "operating_system", + (OperatingSystem.LINUX, OperatingSystem.WINDOWS), +) +def test_get_agent_binary__cached( in_memory_agent_binary_repository: InMemoryAgentBinaryRepository, mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, + operating_system: OperatingSystem, ): - actual_linux_binary = mock_masquerade_agent_binary_repository.get_linux_binary() - in_memory_agent_binary_repository.agent_binaries[OperatingSystem.LINUX] = b"not_linux_binary" - cached_linux_binary = mock_masquerade_agent_binary_repository.get_linux_binary() + actual_linux_binary = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) + in_memory_agent_binary_repository.agent_binaries[operating_system] = b"new_binary" + cached_linux_binary = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) assert actual_linux_binary == cached_linux_binary - - -def test_masquerade_agent_binary_repository_decorator__cached_windows( - in_memory_agent_binary_repository: InMemoryAgentBinaryRepository, - mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, -): - actual_windows_binary = mock_masquerade_agent_binary_repository.get_windows_binary() - in_memory_agent_binary_repository.agent_binaries[ - OperatingSystem.WINDOWS - ] = b"not_windows_binary" - cached_windows_binary = mock_masquerade_agent_binary_repository.get_windows_binary() - - assert actual_windows_binary == cached_windows_binary From ce6d20ba0353ec1e62fd149040d2c3acae11d586 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 11:24:54 -0400 Subject: [PATCH 04/10] UT: Parametrize test_get_agent_binary_retrieval_error() --- .../test_agent_binary_repository.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py index 43f812ab95e..415869e3012 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_agent_binary_repository.py @@ -28,7 +28,7 @@ def agent_binary_repository(): "operating_system,expected_agent_binary", ([OperatingSystem.LINUX, LINUX_AGENT_BINARY], [OperatingSystem.WINDOWS, WINDOWS_AGENT_BINARY]), ) -def test_get_os_binary( +def test_get_agent_binary( agent_binary_repository, operating_system: OperatingSystem, expected_agent_binary: bytes ): assert ( @@ -49,11 +49,9 @@ def error_raising_agent_binary_repository(error_mock_file_repository): return AgentBinaryRepository(error_mock_file_repository) -def test_get_linux_binary_retrieval_error(error_raising_agent_binary_repository): - with pytest.raises(RetrievalError): - error_raising_agent_binary_repository.get_linux_binary() - - -def test_get_windows_binary_retrieval_error(error_raising_agent_binary_repository): +@pytest.mark.parametrize("operating_system", [OperatingSystem.LINUX, OperatingSystem.WINDOWS]) +def test_get_agent_binary_retrieval_error( + error_raising_agent_binary_repository, operating_system: OperatingSystem +): with pytest.raises(RetrievalError): - error_raising_agent_binary_repository.get_windows_binary() + error_raising_agent_binary_repository.get_agent_binary(operating_system) From 89cc05ac40a1d1f6aaea1ec948df087a0d4efdae Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 11:34:50 -0400 Subject: [PATCH 05/10] Island: Use get_agent_binary() in AgentBinaryService --- .../cc/services/agent_binary_service/agent_binary_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py index 1110f2dece9..48e73dbc1f4 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py @@ -20,10 +20,10 @@ def __init__(self, agent_binary_repository: IAgentBinaryRepository): } def get_linux_binary(self) -> BinaryIO: - return self._agent_binary_repository.get_linux_binary() + return self._agent_binary_repository.get_agent_binary(OperatingSystem.LINUX) def get_windows_binary(self) -> BinaryIO: - return self._agent_binary_repository.get_windows_binary() + return self._agent_binary_repository.get_agent_binary(OperatingSystem.WINDOWS) def get_masque(self, operating_system: OperatingSystem) -> Optional[bytes]: return self._os_masque.get(operating_system, None) From b931eb860e3c1b5d70ab1f5e284622e5f7a473e4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 11:42:42 -0400 Subject: [PATCH 06/10] Island: Remove IAgentBinaryRepository.get_{linux_windows}_binary() These were replaced by IAgentBinaryRepository.get_agent_binary(OperatingSystem) --- .../agent_binary_repository.py | 6 ------ .../cc/services/agent_binary_service/build.py | 8 ++++---- .../i_agent_binary_repository.py | 18 ------------------ ...uerade_agent_binary_repository_decorator.py | 10 ---------- .../in_memory_agent_binary_repository.py | 6 ------ 5 files changed, 4 insertions(+), 44 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py index 084e1fe101b..8431f615af9 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_repository.py @@ -21,12 +21,6 @@ def __init__(self, file_repository: IFileRepository): def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: return self._get_binary(AGENT_FILE_NAMES[operating_system]) - def get_linux_binary(self) -> BinaryIO: - return self._get_binary(LINUX_AGENT_FILE_NAME) - - def get_windows_binary(self) -> BinaryIO: - return self._get_binary(WINDOWS_AGENT_FILE_NAME) - def _get_binary(self, filename: str) -> BinaryIO: try: agent_binary = self._file_repository.open_file(filename) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/build.py b/monkey/monkey_island/cc/services/agent_binary_service/build.py index 18c02404248..24311b57ddc 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/build.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/build.py @@ -1,6 +1,7 @@ import logging from pathlib import Path +from common import OperatingSystem from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repositories import ( FileRepositoryCachingDecorator, @@ -49,14 +50,13 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository): :param agent_binary_repository: Used to retrieve the agent binaries """ agent_binaries = { - "Linux": agent_binary_repository.get_linux_binary, - "Windows": agent_binary_repository.get_windows_binary, + "Linux": agent_binary_repository.get_agent_binary(OperatingSystem.LINUX), + "Windows": agent_binary_repository.get_agent_binary(OperatingSystem.WINDOWS), } agent_hashes = {} - for os, get_agent_binary in agent_binaries.items(): + for os, agent_binary in agent_binaries.items(): try: - agent_binary = get_agent_binary() binary_sha256_hash = get_binary_io_sha256_hash(agent_binary) agent_hashes[os] = binary_sha256_hash except RetrievalError as err: diff --git a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py index a75fbe34310..dc0870b145a 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_repository.py @@ -18,21 +18,3 @@ def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: :return: A file-like object that represents the agent binary :raises RetrievalError: If the agent binary could not be retrieved """ - - @abc.abstractmethod - def get_linux_binary(self) -> BinaryIO: - """ - Retrieve linux agent binary - - :return: A file-like object that represents the linux agent binary - :raises RetrievalError: If the agent binary could not be retrieved - """ - - @abc.abstractmethod - def get_windows_binary(self) -> BinaryIO: - """ - Retrieve windows agent binary - - :return: A file-like object that represents the windows agent binary - :raises RetrievalError: If the agent binary could not be retrieved - """ diff --git a/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py b/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py index 43142aaf8eb..b20e0f65147 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py @@ -34,16 +34,6 @@ def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: agent_binary = self._agent_binary_repository.get_agent_binary(operating_system) return self._apply_masque(agent_binary) - @lru_cache() - def get_linux_binary(self) -> BinaryIO: - agent_linux_binary = self._agent_binary_repository.get_linux_binary() - return self._apply_masque(agent_linux_binary) - - @lru_cache() - def get_windows_binary(self) -> BinaryIO: - agent_windows_binary = self._agent_binary_repository.get_windows_binary() - return self._apply_masque(agent_windows_binary) - def _apply_masque(self, agent_binary: BinaryIO) -> BinaryIO: # Note: These null bytes separate the Agent binary from the masque. The goal is to prevent # the masque from being interpreted by the OS as code that should be run. diff --git a/monkey/tests/monkey_island/in_memory_agent_binary_repository.py b/monkey/tests/monkey_island/in_memory_agent_binary_repository.py index 00232903c65..f88bebfdc0a 100644 --- a/monkey/tests/monkey_island/in_memory_agent_binary_repository.py +++ b/monkey/tests/monkey_island/in_memory_agent_binary_repository.py @@ -18,9 +18,3 @@ def __init__(self): def get_agent_binary(self, operating_system: OperatingSystem): return io.BytesIO(self.agent_binaries[operating_system]) - - def get_linux_binary(self): - return io.BytesIO(self.agent_binaries[OperatingSystem.LINUX]) - - def get_windows_binary(self): - return io.BytesIO(self.agent_binaries[OperatingSystem.WINDOWS]) From 354800a4ed5ece076502379c9a3de9a6a1ec5a6a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 12:03:05 -0400 Subject: [PATCH 07/10] Island: Add IAgentBinaryService.get_agent_binary() --- .../agent_binary_service/agent_binary_service.py | 3 +++ .../agent_binary_service/i_agent_binary_service.py | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py index 48e73dbc1f4..5d4089cf247 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py @@ -19,6 +19,9 @@ def __init__(self, agent_binary_repository: IAgentBinaryRepository): OperatingSystem.WINDOWS: None, } + def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + return self._agent_binary_repository.get_agent_binary(operating_system) + def get_linux_binary(self) -> BinaryIO: return self._agent_binary_repository.get_agent_binary(OperatingSystem.LINUX) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py index f690ab4209e..f004407a700 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py @@ -9,6 +9,16 @@ class IAgentBinaryService(metaclass=abc.ABCMeta): A service for retrieving and manipulating agent binaries """ + @abc.abstractmethod + def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + """ + Retrieve an agent binary for the specified operating system + + :param operating_system: The operating system with which the binary must be compatible + :return: A file-like object that represents the agent binary + :raises RetrievalError: If the agent binary could not be retrieved + """ + @abc.abstractmethod def get_linux_binary(self) -> BinaryIO: """ From 1d78e322261f8a8ddf86703360463674037adab8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 12:03:33 -0400 Subject: [PATCH 08/10] Island: Use get_agent_binary() in LocalMonkeyRunService --- .../monkey_island/cc/services/run_local_monkey.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/services/run_local_monkey.py b/monkey/monkey_island/cc/services/run_local_monkey.py index 9f34ca4af5d..58b9637cc00 100644 --- a/monkey/monkey_island/cc/services/run_local_monkey.py +++ b/monkey/monkey_island/cc/services/run_local_monkey.py @@ -8,6 +8,7 @@ from shutil import copyfileobj from typing import Sequence +from common import OperatingSystem from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE from common.types import OTP from monkey_island.cc.repositories import RetrievalError @@ -16,7 +17,10 @@ logger = logging.getLogger(__name__) -AGENT_NAMES = {"linux": "monkey-linux-64", "windows": "monkey-windows-64.exe"} +AGENT_NAMES = { + OperatingSystem.LINUX: "monkey-linux-64", + OperatingSystem.WINDOWS: "monkey-windows-64.exe", +} class LocalMonkeyRunService: @@ -32,14 +36,9 @@ def __init__( def run_local_monkey(self, otp: OTP): # get the monkey executable suitable to run on the server - operating_system = platform.system().lower() try: - agents = { - "linux": self._agent_binary_service.get_linux_binary, - "windows": self._agent_binary_service.get_windows_binary, - } - - agent_binary = agents[platform.system().lower()]() + operating_system = OperatingSystem[platform.system().upper()] + agent_binary = self._agent_binary_service.get_agent_binary(operating_system) except RetrievalError as err: logger.error( f"No Agent can be retrieved for the specified operating system" From d333553b5b01011b0d09ee8c45c2186e983f4745 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 12:03:56 -0400 Subject: [PATCH 09/10] Island: Use get_agent_binary() in AgentBinaries resource --- .../flask_resources/agent_binaries.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/flask_resources/agent_binaries.py b/monkey/monkey_island/cc/services/agent_binary_service/flask_resources/agent_binaries.py index df159471655..339b72e28fd 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/flask_resources/agent_binaries.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/flask_resources/agent_binaries.py @@ -3,6 +3,7 @@ from flask import make_response, send_file +from common import OperatingSystem from monkey_island.cc.flask_utils import AbstractResource from monkey_island.cc.repositories import RetrievalError @@ -26,12 +27,8 @@ def get(self, os): :return: an agent binary file """ try: - agent_binaries = { - "linux": self._agent_binary_service.get_linux_binary, - "windows": self._agent_binary_service.get_windows_binary, - } - - file = agent_binaries[os]() + operating_system = OperatingSystem[os.upper()] + file = self._agent_binary_service.get_agent_binary(operating_system) return send_file(file, mimetype="application/octet-stream") except KeyError as err: From 7d984c206cdc1cd9627e7c1e78e999b0e146d0e0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 25 Apr 2023 12:04:45 -0400 Subject: [PATCH 10/10] Island: Remove disused IAgentBinaryService.get{linux,windows}_binaries --- .../agent_binary_service.py | 6 ------ .../i_agent_binary_service.py | 18 ------------------ 2 files changed, 24 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py index 5d4089cf247..4f85a43b555 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/agent_binary_service.py @@ -22,12 +22,6 @@ def __init__(self, agent_binary_repository: IAgentBinaryRepository): def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: return self._agent_binary_repository.get_agent_binary(operating_system) - def get_linux_binary(self) -> BinaryIO: - return self._agent_binary_repository.get_agent_binary(OperatingSystem.LINUX) - - def get_windows_binary(self) -> BinaryIO: - return self._agent_binary_repository.get_agent_binary(OperatingSystem.WINDOWS) - def get_masque(self, operating_system: OperatingSystem) -> Optional[bytes]: return self._os_masque.get(operating_system, None) diff --git a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py index f004407a700..e4d953efad8 100644 --- a/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py +++ b/monkey/monkey_island/cc/services/agent_binary_service/i_agent_binary_service.py @@ -19,24 +19,6 @@ def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: :raises RetrievalError: If the agent binary could not be retrieved """ - @abc.abstractmethod - def get_linux_binary(self) -> BinaryIO: - """ - Retrieve linux agent binary - - :return: A file-like object that represents the linux agent binary - :raises RetrievalError: If the agent binary could not be retrieved - """ - - @abc.abstractmethod - def get_windows_binary(self) -> BinaryIO: - """ - Retrieve windows agent binary - - :return: A file-like object that represents the windows agent binary - :raises RetrievalError: If the agent binary could not be retrieved - """ - @abc.abstractmethod def get_masque(self, operating_system: OperatingSystem) -> Optional[bytes]: """