Merge branch '3309-fix-masquerade-caching' into release/2.2.0

Issue #3309
PR #3387

envs/monkey_zoo/blackbox/island_client/monkey_island_client.py

@@ -40,11 +40,16 @@ def set_masque(self, masque):
         masque = b"" if masque is None else masque
         for operating_system in [operating_system.name for operating_system in OperatingSystem]:
             if self.requests.put(f"api/agent-binaries/{operating_system}/masque", data=masque).ok:
-                LOGGER.info(f'Setting {operating_system} masque to "{masque}"')
+                formatted_masque = masque if len(masque) <= 64 else (masque[:64] + b"...")
+                LOGGER.info(f'Setting {operating_system} masque to "{formatted_masque}"')
             else:
                 LOGGER.error(f"Failed to set {operating_system} masque")
                 assert False
 
+    def get_agent_binary(self, operating_system: OperatingSystem) -> bytes:
+        response = self.requests.get(f"api/agent-binaries/{operating_system.name}")
+        return response.content
+
     def get_propagation_credentials(self) -> Sequence[Credentials]:
         response = self.requests.get("api/propagation-credentials")
         return [Credentials(**credentials) for credentials in response.json()]

envs/monkey_zoo/blackbox/test_blackbox.py

@@ -1,15 +1,17 @@
 import logging
 import os
+import random
 from copy import deepcopy
 from http import HTTPStatus
 from threading import Thread
 from time import sleep
-from typing import List, Optional, Sequence
+from typing import Iterable, List, Optional, Sequence
 from uuid import uuid4
 
 import pytest
 import requests
 
+from common import OperatingSystem
 from common.credentials import Credentials, NTHash, Password, Username
 from common.types import OTP, SocketAddress
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
@@ -480,6 +482,38 @@ def test_agent_logout(island):
     assert agent_requests.get(GET_AGENT_OTP_ENDPOINT).status_code == HTTPStatus.UNAUTHORIZED
 
 
+RANDBYTES_SIZE = 1024 * 1024 * 2  # 2MB
+LINUX_DATA = [
+    b"This is a string for Linux!",
+    random.randbytes(RANDBYTES_SIZE),  # noqa: DUO102
+    b"More strings",
+    b"A much longer supercalifragilisticexpialidocious string to be included in the masque.",
+]
+WINDOWS_DATA = [
+    b"This is a string for Windows!",
+    random.randbytes(RANDBYTES_SIZE),  # noqa: DUO102
+    LINUX_DATA[2],
+    LINUX_DATA[3],
+]
+
+
+@pytest.mark.parametrize(
+    "data, operating_system",
+    [(LINUX_DATA, OperatingSystem.LINUX), (WINDOWS_DATA, OperatingSystem.WINDOWS)],
+    ids=["Linux", "Windows"],
+)
+def test_masquerade(
+    island_client: MonkeyIslandClient, data: Iterable[bytes], operating_system: OperatingSystem
+):
+    masque_data = b"\0".join(data)
+
+    island_client.set_masque(masque_data)
+    agent_binary = island_client.get_agent_binary(operating_system)
+
+    for d in data:
+        assert d in agent_binary
+
+
 # NOTE: These test methods are ordered to give time for the slower zoo machines
 # to boot up and finish starting services.
 # noinspection PyUnresolvedReferences

monkey/monkey_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py

@@ -1,4 +1,3 @@
-from functools import lru_cache
 from typing import BinaryIO
 
 from common import OperatingSystem
@@ -35,7 +34,6 @@ def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO:
 
         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)

monkey/tests/unit_tests/monkey_island/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py

@@ -38,7 +38,7 @@ def in_memory_agent_binary_repository() -> InMemoryAgentBinaryRepository:
 
 
 @pytest.fixture
-def in_memory_masquerade_repository() -> InMemoryAgentBinaryRepository:
+def in_memory_masquerade_repository() -> IMasqueradeRepository:
     return InMemoryMasqueradeRepository()
 
 
@@ -72,40 +72,6 @@ def test_get_agent_binary(
     assert actual_binary.read() == expected_agent_binary.getvalue()
 
 
-@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_binary = mock_masquerade_agent_binary_repository.get_agent_binary(operating_system)
-    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.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()
-
-
 def test_one_unset_masque(
     in_memory_agent_binary_repository: IAgentBinaryRepository,
     in_memory_masquerade_repository: IMasqueradeRepository,