3241 refactor get agent binary methods

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

@@ -1,22 +1,25 @@
 from typing import BinaryIO
 
+from common import OperatingSystem
 from monkey_island.cc.repositories import IFileRepository, RetrievalError
 
 from .i_agent_binary_repository import IAgentBinaryRepository
 
 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_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_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO:
+        return self._get_binary(AGENT_FILE_NAMES[operating_system])
 
     def _get_binary(self, filename: str) -> BinaryIO:
         try:

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

@@ -19,11 +19,8 @@ def __init__(self, agent_binary_repository: IAgentBinaryRepository):
             OperatingSystem.WINDOWS: None,
         }
 
-    def get_linux_binary(self) -> BinaryIO:
-        return self._agent_binary_repository.get_linux_binary()
-
-    def get_windows_binary(self) -> BinaryIO:
-        return self._agent_binary_repository.get_windows_binary()
+    def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO:
+        return self._agent_binary_repository.get_agent_binary(operating_system)
 
     def get_masque(self, operating_system: OperatingSystem) -> Optional[bytes]:
         return self._os_masque.get(operating_system, None)

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:

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:

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

@@ -1,26 +1,20 @@
 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_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:
+    def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO:
         """
-        Retrieve windows agent binary
+        Retrieve an agent binary for the specified operating system
 
-        :return: A file-like object that represents the windows agent binary
+        :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
         """

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

@@ -10,20 +10,12 @@ class IAgentBinaryService(metaclass=abc.ABCMeta):
     """
 
     @abc.abstractmethod
-    def get_linux_binary(self) -> BinaryIO:
+    def get_agent_binary(self, operating_system: OperatingSystem) -> BinaryIO:
         """
-        Retrieve linux agent binary
+        Retrieve an agent binary for the specified operating system
 
-        :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
+        :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
         """
 

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
@@ -28,14 +30,9 @@ def __init__(
         self._null_bytes_length = null_bytes_length
 
     @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 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)
 
     def _apply_masque(self, agent_binary: BinaryIO) -> BinaryIO:
         # Note: These null bytes separate the Agent binary from the masque. The goal is to prevent

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"

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

monkey/tests/monkey_island/in_memory_agent_binary_repository.py

@@ -0,0 +1,20 @@
+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])

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

@@ -1,16 +1,41 @@
+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,
 )
 
 LINUX_AGENT_BINARY = b"linux_binary"
 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_agent_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,13 @@ 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):
-    with pytest.raises(RetrievalError):
-        agent_binary_repository.get_linux_binary()
-
-
-def test_get_windows_binary_retrieval_error(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):
-        agent_binary_repository.get_windows_binary()
+        error_raising_agent_binary_repository.get_agent_binary(operating_system)