diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 88236f771..885d467be 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,2 +1,3 @@ from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository -from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrivalError +from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError +from .agent_binary_repository import AgentBinaryRepository diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py index fc1ac315e..4c1334f71 100644 --- a/monkey/monkey_island/cc/repository/agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py @@ -1,6 +1,6 @@ from typing import BinaryIO -from . import AgentRetrivalError, FileRetrivalError, IAgentBinaryRepository, IFileRepository +from . import AgentRetrievalError, FileRetrievalError, IAgentBinaryRepository, IFileRepository LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" @@ -10,18 +10,18 @@ class AgentBinaryRepository(IAgentBinaryRepository): def __init__(self, file_repository: IFileRepository): self._file_repository = file_repository - def __get_binary(self, filename) -> BinaryIO: + 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_binary(self, filename: str) -> BinaryIO: try: agent_binary = self._file_repository.open_file(filename) return agent_binary - except FileRetrivalError as err: - raise AgentRetrivalError( + except FileRetrievalError as err: + raise AgentRetrievalError( f"An error occurred while retrieving the {filename}" f" agent binary from {self._file_repository}: {err}" ) - - def get_linux_binary(self) -> BinaryIO: - self.__get_binary(LINUX_AGENT_FILE_NAME) - - def get_windows_binary(self) -> BinaryIO: - self.__get_binary(WINDOWS_AGENT_FILE_NAME) 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 aac33176c..56cf2f96b 100644 --- a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py @@ -2,7 +2,7 @@ import abc from typing import BinaryIO -class AgentRetrivalError(IOError): +class AgentRetrievalError(IOError): pass