From b9efffd1c2e16a7db2a760316b4a26c48e74e4a1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 26 Apr 2023 09:30:29 -0400 Subject: [PATCH] Island: Use make_fileobj_copy() in MasqueradeAgent...Decorator --- ...erade_agent_binary_repository_decorator.py | 8 +++++++- ...erade_agent_binary_repository_decorator.py | 20 ++++++++++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) 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 3228fedbcc0..ae827cbbbcd 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 @@ -3,6 +3,7 @@ from typing import BinaryIO, Mapping, Optional from common import OperatingSystem +from common.utils.file_utils import make_fileobj_copy from .i_agent_binary_repository import IAgentBinaryRepository @@ -29,8 +30,13 @@ def __init__( self._masques = masques self._null_bytes = b"\x00" * null_bytes_length - @lru_cache() def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO: + original_file = self._get_agent_binary(operating_system) + + return make_fileobj_copy(original_file) + + @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(operating_system, agent_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 8d5bdef8309..43ba6f7579d 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 @@ -74,7 +74,25 @@ def test_get_agent_binary__cached( in_memory_agent_binary_repository.agent_binaries[operating_system] = b"new_binary" cached_binary = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) - assert actual_binary == cached_binary + assert actual_binary.read() == cached_binary.read() + + +def test_get_agent_binary__cached_multiple_calls( + in_memory_agent_binary_repository: InMemoryAgentBinaryRepository, + mock_masquerade_agent_binary_repository: MasqueradeAgentBinaryRepositoryDecorator, +): + operating_system = OperatingSystem.WINDOWS + + cached_binary_1 = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) + in_memory_agent_binary_repository.agent_binaries[operating_system] = b"new_binary" + cached_binary_2 = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) + cached_binary_3 = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system) + + # Writing the assertion this way verifies that returned files have had their positions reset to + # the beginning (i.e. seek(0)). + assert cached_binary_1.read() == MASQUED_WINDOWS_AGENT_BINARY.getvalue() + assert cached_binary_2.read() == MASQUED_WINDOWS_AGENT_BINARY.getvalue() + assert cached_binary_3.read() == MASQUED_WINDOWS_AGENT_BINARY.getvalue() @pytest.mark.parametrize(