From 47df2575459f2e9a1392d6f8d0f44dbf97ae1ee4 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
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}"'