diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index e3cf5c2fc..88236f771 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,2 +1,2 @@ from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository -from .i_agent_binary_repository import IAgentBinaryRepository +from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrivalError diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py new file mode 100644 index 000000000..ec414c95c --- /dev/null +++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py @@ -0,0 +1,29 @@ +from typing import BinaryIO + +from . import AgentRetrivalError, FileRetrivalError, IAgentBinaryRepository, IFileRepository + +LINUX_AGENT_FILE_NAME = "monkey-linux-64" +WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" + + +class AgentBinaryRepository(IAgentBinaryRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + + def get_linux_binary(self) -> BinaryIO: + try: + agent_binary = self._file_repository.open_file(LINUX_AGENT_FILE_NAME) + return agent_binary + except FileRetrivalError as err: + raise AgentRetrivalError( + f"An error occurred while retrieving the Linux agent binary: {err}" + ) + + def get_windows_binary(self) -> BinaryIO: + try: + agent_binary = self._file_repository.open_file(WINDOWS_AGENT_FILE_NAME) + return agent_binary + except FileRetrivalError as err: + raise AgentRetrivalError( + f"An error occurred while retrieving the Windows agent binary: {err}" + ) diff --git a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py index a1d858e05..aac33176c 100644 --- a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py @@ -2,6 +2,10 @@ import abc from typing import BinaryIO +class AgentRetrivalError(IOError): + pass + + class IAgentBinaryRepository(metaclass=abc.ABCMeta): """ A repository that retrieves the agent binaries diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 431eb68d3..40ff4e69b 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -2,7 +2,13 @@ from pathlib import Path from common import DIContainer from common.aws import AWSInstance -from monkey_island.cc.repository import IFileRepository, LocalStorageFileRepository +from monkey_island.cc.repository import ( + AgentBinaryRepository, + IAgentBinaryRepository, + IFileRepository, + LocalStorageFileRepository, +) +from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.services import AWSService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -10,6 +16,8 @@ from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService from . import AuthenticationService, JsonFileUserDatastore from .reporting.report import ReportService +AGENT_BINARIES_PATH = Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries" + def initialize_services(data_dir: Path) -> DIContainer: container = DIContainer() @@ -19,6 +27,7 @@ def initialize_services(data_dir: Path) -> DIContainer: IFileRepository, LocalStorageFileRepository(data_dir / "custom_pbas") ) container.register_instance(AWSService, container.resolve(AWSService)) + container.register_instance(IAgentBinaryRepository, _build_agent_binary_repository()) # This is temporary until we get DI all worked out. PostBreachFilesService.initialize(container.resolve(IFileRepository)) @@ -27,3 +36,8 @@ def initialize_services(data_dir: Path) -> DIContainer: ReportService.initialize(container.resolve(AWSService)) return container + + +def _build_agent_binary_repository(): + file_repository = LocalStorageFileRepository(AGENT_BINARIES_PATH) + return AgentBinaryRepository(file_repository)