3244 agent binary repository decorator

monkey/common/utils/file_utils/__init__.py

@@ -5,6 +5,7 @@
     get_text_file_contents,
     InvalidPath,
     make_fileobj_copy,
+    append_bytes,
 )
 from .secure_directory import create_secure_directory
 from .secure_file import open_new_securely_permissioned_file

monkey/common/utils/file_utils/file_utils.py

@@ -66,3 +66,13 @@ def make_fileobj_copy(src: BinaryIO) -> BinaryIO:
     dst.seek(0)
 
     return dst
+
+
+def append_bytes(file: BinaryIO, bytes_to_append: bytes) -> BinaryIO:
+    starting_position = file.tell()
+
+    file.seek(0, io.SEEK_END)
+    file.write(bytes_to_append)
+    file.seek(starting_position, io.SEEK_SET)
+
+    return file

monkey/infection_monkey/exploit/__init__.py

@@ -1,5 +1,6 @@
 from .i_agent_binary_repository import IAgentBinaryRepository, RetrievalError
 from .caching_agent_binary_repository import CachingAgentBinaryRepository
+from .polymorphic_agent_binary_repository_decorator import PolymorphicAgentBinaryRepositoryDecorator
 from .island_api_agent_otp_provider import IslandAPIAgentOTPProvider
 from .i_agent_otp_provider import IAgentOTPProvider
 from .exploiter_wrapper import ExploiterWrapper

monkey/infection_monkey/exploit/polymorphic_agent_binary_repository_decorator.py

@@ -0,0 +1,37 @@
+import io
+from random import randbytes  # noqa: DUO102
+
+from common import OperatingSystem
+from common.utils.file_utils import append_bytes
+
+from . import IAgentBinaryRepository
+
+DEFAULT_NULL_BYTES_LENGTH = 16
+DEFAULT_RANDOM_BYTES_LENGTH = 8
+
+
+class PolymorphicAgentBinaryRepositoryDecorator(IAgentBinaryRepository):
+    """
+    PolymorphicAgentBinaryRepositoryDecorator adds random bytes to agent binaries in a way that does
+    not affect their functionality. This allows Infection Monkey to emulate a property of
+    polymorphic malware: all copies have different hashes.
+    """
+
+    def __init__(
+        self,
+        agent_binary_repository: IAgentBinaryRepository,
+        null_bytes_length: int = DEFAULT_NULL_BYTES_LENGTH,
+        random_bytes_length: int = DEFAULT_RANDOM_BYTES_LENGTH,
+    ):
+        self._agent_binary_repository = agent_binary_repository
+        self._random_bytes_length = random_bytes_length
+        self._null_bytes = b"\x00" * null_bytes_length
+
+    def get_agent_binary(self, operating_system: OperatingSystem) -> io.BytesIO:
+        agent_binary = self._agent_binary_repository.get_agent_binary(operating_system)
+
+        # 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.
+        append_bytes(agent_binary, self._null_bytes + randbytes(self._random_bytes_length))
+
+        return agent_binary

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

@@ -1,9 +1,8 @@
-import io
 from functools import lru_cache
 from typing import BinaryIO
 
 from common import OperatingSystem
-from common.utils.file_utils import make_fileobj_copy
+from common.utils.file_utils import append_bytes, make_fileobj_copy
 
 from .i_agent_binary_repository import IAgentBinaryRepository
 from .i_masquerade_repository import IMasqueradeRepository
@@ -49,8 +48,6 @@ def _apply_masque(self, operating_system: OperatingSystem, agent_binary: BinaryI
 
         # 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.
-        agent_binary.seek(0, io.SEEK_END)
-        agent_binary.write(self._null_bytes + masque)
-        agent_binary.seek(0)
+        append_bytes(agent_binary, self._null_bytes + masque)
 
         return agent_binary

monkey/tests/unit_tests/common/utils/test_file_utils.py

@@ -7,6 +7,7 @@
 
 from common.utils.environment import is_windows_os
 from common.utils.file_utils import (
+    append_bytes,
     create_secure_directory,
     make_fileobj_copy,
     open_new_securely_permissioned_file,
@@ -159,3 +160,22 @@ def test_make_fileobj_copy_seek_src_to_0():
         # their positions reset to 0.
         assert src.read() == TEST_STR
         assert dst.read() == TEST_STR
+
+
+def test_append_bytes__pos_0():
+    bytes_io = io.BytesIO(b"1234 5678")
+
+    append_bytes(bytes_io, b"abcd")
+
+    assert bytes_io.read() == b"1234 5678abcd"
+
+
+def test_append_bytes__pos_5():
+    bytes_io = io.BytesIO(b"1234 5678")
+    bytes_io.seek(5, io.SEEK_SET)
+
+    append_bytes(bytes_io, b"abcd")
+
+    assert bytes_io.read() == b"5678abcd"
+    bytes_io.seek(0, io.SEEK_SET)
+    assert bytes_io.read() == b"1234 5678abcd"

monkey/tests/unit_tests/infection_monkey/exploit/test_polymorphic_agent_binary_repository_decorator.py

@@ -0,0 +1,59 @@
+import io
+from unittest.mock import MagicMock
+
+import pytest
+
+from common import OperatingSystem
+from infection_monkey.exploit import (
+    IAgentBinaryRepository,
+    PolymorphicAgentBinaryRepositoryDecorator,
+)
+
+AGENT_BINARY = b"Agent Binary Bytes!"
+NULL_BYTES_LENGTH = 4
+RANDOM_BYTES_LENGTH = 4
+
+
+@pytest.fixture
+def mock_agent_binary_repository() -> IAgentBinaryRepository:
+    mabr = MagicMock(IAgentBinaryRepository)
+    mabr.get_agent_binary.return_value = io.BytesIO(AGENT_BINARY)
+
+    return mabr
+
+
+@pytest.fixture
+def polymorphic_agent_binary_repository(
+    mock_agent_binary_repository: IAgentBinaryRepository,
+) -> IAgentBinaryRepository:
+    return PolymorphicAgentBinaryRepositoryDecorator(
+        mock_agent_binary_repository, NULL_BYTES_LENGTH, RANDOM_BYTES_LENGTH
+    )
+
+
+def test_get_agent_binary(polymorphic_agent_binary_repository: IAgentBinaryRepository):
+    null_bytes_start_index = len(AGENT_BINARY)
+    null_bytes_stop_index = len(AGENT_BINARY) + NULL_BYTES_LENGTH
+
+    agent_binary = polymorphic_agent_binary_repository.get_agent_binary(OperatingSystem.LINUX)
+    agent_binary_bytes = agent_binary.read()
+
+    assert agent_binary_bytes[:null_bytes_start_index] == AGENT_BINARY
+    assert (
+        agent_binary_bytes[null_bytes_start_index:null_bytes_stop_index]
+        == b"\x00" * NULL_BYTES_LENGTH
+    )
+
+
+def test_get_agent_binary__random(polymorphic_agent_binary_repository: IAgentBinaryRepository):
+    random_bytes_start_index = len(AGENT_BINARY) + NULL_BYTES_LENGTH
+
+    agent_binary = polymorphic_agent_binary_repository.get_agent_binary(OperatingSystem.LINUX)
+    agent_binary_bytes_1 = agent_binary.read()
+    agent_binary = polymorphic_agent_binary_repository.get_agent_binary(OperatingSystem.LINUX)
+    agent_binary_bytes_2 = agent_binary.read()
+
+    assert (
+        agent_binary_bytes_1[random_bytes_start_index:]
+        != agent_binary_bytes_2[random_bytes_start_index:]
+    )