Island: Raise RetrievalError from IAgentBinaryRepository

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

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 IAgentBinaryRepository, IFileRepository, RetrievalError
-
-from . import AgentRetrievalError, IAgentBinaryRepository, IFileRepository
 
 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 Exception as err:
-        except repository.FileNotFoundError as err:
+            raise RetrievalError(
-            raise AgentRetrievalError(
                 f"An error occurred while retrieving the {filename}"
                 f" agent binary from {self._file_repository}: {err}"
             )

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
         """

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)

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():

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}"'