From 47df2575459f2e9a1392d6f8d0f44dbf97ae1ee4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:18:42 -0400 Subject: [PATCH] Island: Raise RetrievalError from IAgentBinaryRepository --- monkey/monkey_island/cc/repository/__init__.py | 2 +- .../cc/repository/agent_binary_repository.py | 9 +++------ .../cc/repository/i_agent_binary_repository.py | 6 ++---- monkey/monkey_island/cc/resources/agent_binaries.py | 8 ++++---- monkey/monkey_island/cc/services/initialize.py | 4 ++-- monkey/monkey_island/cc/services/run_local_monkey.py | 4 ++-- 6 files changed, 14 insertions(+), 19 deletions(-) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index a44e7b1ad..8fadb9ccf 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,6 +1,6 @@ from .errors import RetrievalError from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository -from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError +from .i_agent_binary_repository import IAgentBinaryRepository from .agent_binary_repository import AgentBinaryRepository from .i_agent_configuration_repository import IAgentConfigurationRepository from .file_agent_configuration_repository import FileAgentConfigurationRepository diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py index a09d88868..9b753ce00 100644 --- a/monkey/monkey_island/cc/repository/agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py @@ -1,8 +1,6 @@ from typing import BinaryIO -from monkey_island.cc import repository - -from . import AgentRetrievalError, IAgentBinaryRepository, IFileRepository +from . import IAgentBinaryRepository, IFileRepository, RetrievalError LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" @@ -22,9 +20,8 @@ class AgentBinaryRepository(IAgentBinaryRepository): try: agent_binary = self._file_repository.open_file(filename) return agent_binary - # TODO: Reevaluate this - except repository.FileNotFoundError as err: - raise AgentRetrievalError( + except Exception as err: + raise RetrievalError( f"An error occurred while retrieving the {filename}" f" agent binary from {self._file_repository}: {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 56cf2f96b..6c9abcc40 100644 --- a/monkey/monkey_island/cc/repository/i_agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_binary_repository.py @@ -2,10 +2,6 @@ import abc from typing import BinaryIO -class AgentRetrievalError(IOError): - pass - - class IAgentBinaryRepository(metaclass=abc.ABCMeta): """ A repository that retrieves the agent binaries @@ -17,6 +13,7 @@ class IAgentBinaryRepository(metaclass=abc.ABCMeta): 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 @@ -25,4 +22,5 @@ class IAgentBinaryRepository(metaclass=abc.ABCMeta): Retrieve windows agent binary :return: A file-like object that represents the windows agent binary + :raises RetrievalError: If the agent binary could not be retrieved """ diff --git a/monkey/monkey_island/cc/resources/agent_binaries.py b/monkey/monkey_island/cc/resources/agent_binaries.py index 8d960b6a4..0d746b932 100644 --- a/monkey/monkey_island/cc/resources/agent_binaries.py +++ b/monkey/monkey_island/cc/resources/agent_binaries.py @@ -2,7 +2,7 @@ import logging from flask import make_response, send_file -from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository +from monkey_island.cc.repository import IAgentBinaryRepository, RetrievalError from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -31,10 +31,10 @@ class AgentBinaries(AbstractResource): file = agent_binaries[os]() return send_file(file, mimetype="application/octet-stream") - except AgentRetrievalError as err: - logger.error(err) - return make_response({"error": str(err)}, 500) except KeyError as err: error_msg = f'No Agents are available for unsupported operating system "{os}": {err}' logger.error(error_msg) return make_response({"error": error_msg}, 404) + except RetrievalError as err: + logger.error(err) + return make_response({"error": str(err)}, 500) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 74d447c90..4b03daf7d 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -6,12 +6,12 @@ from common.aws import AWSInstance from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, - AgentRetrievalError, FileAgentConfigurationRepository, IAgentBinaryRepository, IAgentConfigurationRepository, IFileRepository, LocalStorageFileRepository, + RetrievalError, ) from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.services import AWSService @@ -76,7 +76,7 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository): agent_binary = get_agent_binary() binary_sha256_hash = get_binary_io_sha256_hash(agent_binary) agent_hashes[os] = binary_sha256_hash - except AgentRetrievalError as err: + except RetrievalError as err: logger.error(f"No agent available for {os}: {err}") for os, binary_sha256_hash in agent_hashes.items(): diff --git a/monkey/monkey_island/cc/services/run_local_monkey.py b/monkey/monkey_island/cc/services/run_local_monkey.py index a54e2c6b9..a56642e2c 100644 --- a/monkey/monkey_island/cc/services/run_local_monkey.py +++ b/monkey/monkey_island/cc/services/run_local_monkey.py @@ -5,7 +5,7 @@ import subprocess from pathlib import Path from shutil import copyfileobj -from monkey_island.cc.repository import AgentRetrievalError, IAgentBinaryRepository +from monkey_island.cc.repository import IAgentBinaryRepository, RetrievalError from monkey_island.cc.server_utils.consts import ISLAND_PORT from monkey_island.cc.services.utils.network_utils import local_ip_addresses @@ -29,7 +29,7 @@ class LocalMonkeyRunService: } agent_binary = agents[platform.system().lower()]() - except AgentRetrievalError as err: + except RetrievalError as err: logger.error( f"No Agent can be retrieved for the specified operating system" f'"{operating_system}"'