Island: Use make_fileobj_copy() in MasqueradeAgent...Decorator

guardicore · Apr 26, 2023 · b9efffd · b9efffd
1 parent bcbb52c
commit b9efffd
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/...y_island/cc/services/agent_binary_service/masquerade_agent_binary_repository_decorator.py b/...y_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/...and/cc/services/agent_binary_service/test_masquerade_agent_binary_repository_decorator.py b/...and/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(