diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index cc89c5c57..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 FileRetrievalError, IFileRepository, LocalStorageFileRepository -from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError +from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository +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 4c1334f71..9b753ce00 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 AgentRetrievalError, FileRetrievalError, IAgentBinaryRepository, IFileRepository +from . import IAgentBinaryRepository, IFileRepository, RetrievalError LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" @@ -20,8 +20,8 @@ class AgentBinaryRepository(IAgentBinaryRepository): try: agent_binary = self._file_repository.open_file(filename) return agent_binary - except FileRetrievalError 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/file_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py index fcccd49a1..89fc57960 100644 --- a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -5,6 +5,7 @@ from common.configuration import ( AgentConfiguration, AgentConfigurationSchema, ) +from monkey_island.cc import repository from monkey_island.cc.repository import ( IAgentConfigurationRepository, IFileRepository, @@ -25,10 +26,10 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository): configuration_json = f.read().decode() return self._schema.loads(configuration_json) - # TODO: Handle FileRetrievalError vs FileNotFoundError - # https://github.com/guardicore/monkey/blob/e8001d8cf76340e42bf17ff62523bd2d85fc4841/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py#L47-L50 - except RetrievalError: + except repository.FileNotFoundError: return self._schema.loads(DEFAULT_AGENT_CONFIGURATION) + except Exception as err: + raise RetrievalError(f"Error retrieving the agent configuration: {err}") def store_configuration(self, agent_configuration: AgentConfiguration): configuration_json = self._schema.dumps(agent_configuration) diff --git a/monkey/monkey_island/cc/repository/file_storage/__init__.py b/monkey/monkey_island/cc/repository/file_storage/__init__.py index 30b59b5d9..14227bc66 100644 --- a/monkey/monkey_island/cc/repository/file_storage/__init__.py +++ b/monkey/monkey_island/cc/repository/file_storage/__init__.py @@ -1,2 +1,2 @@ -from .i_file_repository import IFileRepository, FileRetrievalError +from .i_file_repository import IFileRepository, FileNotFoundError from .local_storage_file_repository import LocalStorageFileRepository diff --git a/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py b/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py index 2531fd711..bb9e345dc 100644 --- a/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/file_storage/i_file_repository.py @@ -4,7 +4,7 @@ from typing import BinaryIO from monkey_island.cc.repository import RetrievalError -class FileRetrievalError(RetrievalError): +class FileNotFoundError(RetrievalError): pass @@ -30,7 +30,8 @@ class IFileRepository(metaclass=abc.ABCMeta): :param unsafe_file_name: An unsanitized file name that identifies the file to be opened :return: A file-like object providing access to the file's contents - :raises FileRetrievalError: if the file cannot be opened + :raises FileNotFoundError: if the file does not exist + :raises RetrievalError: if the file exists but cannot be retrieved """ pass diff --git a/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py b/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py index 4202496ad..cb18c6ba8 100644 --- a/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py @@ -4,9 +4,10 @@ from pathlib import Path from typing import BinaryIO from common.utils.file_utils import get_all_regular_files_in_directory +from monkey_island.cc.repository import RetrievalError from monkey_island.cc.server_utils.file_utils import create_secure_directory -from . import FileRetrievalError, IFileRepository +from . import IFileRepository, i_file_repository logger = logging.getLogger(__name__) @@ -43,12 +44,15 @@ class LocalStorageFileRepository(IFileRepository): try: logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") - except OSError as err: - # TODO: The interface should make a destinction between file not found and an error when - # retrieving a file that should exist. The built-in `FileNotFoundError` is not - # sufficient because it inherits from `OSError` and the interface does not - # guarantee that the file is stored on the local file system. - raise FileRetrievalError(f"Failed to retrieve file {safe_file_path}: {err}") from err + except FileNotFoundError as err: + # Wrap Python's FileNotFound error, which is-an OSError, in repository.FileNotFoundError + raise i_file_repository.FileNotFoundError( + f'The requested file "{unsafe_file_name}" does not exist: {err}' + ) + except Exception as err: + raise RetrievalError( + f'Error retrieving file "{unsafe_file_name}" from the repository: {err}' + ) def delete_file(self, unsafe_file_name: str): safe_file_path = self._get_safe_file_path(unsafe_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 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/repository/i_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py index 105d26fe9..1e85fadf5 100644 --- a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py @@ -14,8 +14,8 @@ class IAgentConfigurationRepository(ABC): Retrieve the agent configuration from the repository :return: The agent configuration as retrieved from the repository, or the default - configuration if none could be retrieved. - :raises RetrievalError: if the configuration can not be retrieved + configuration if the repository is empty + :raises RetrievalError: if the configuration could not be retrieved """ pass 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/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index 85395d5e3..d8359049a 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -2,7 +2,8 @@ import logging from flask import make_response, send_file -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__file__) @@ -24,7 +25,6 @@ class PBAFileDownload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except FileRetrievalError as err: - error_msg = f"Failed to open file {filename}: {err}" - logger.error(error_msg) - return make_response({"error": error_msg}, 404) + except repository.FileNotFoundError as err: + logger.error(str(err)) + return make_response({"error": str(err)}, 404) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 9d5cce8b1..277804a72 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -5,7 +5,8 @@ from flask import Response, make_response, request, send_file from werkzeug.utils import secure_filename as sanitize_filename from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.config import ConfigService @@ -57,10 +58,9 @@ class FileUpload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except FileRetrievalError as err: - error_msg = f"Failed to open file {filename}: {err}" - logger.error(error_msg) - return make_response({"error": error_msg}, 404) + except repository.FileNotFoundError as err: + logger.error(str(err)) + return make_response({"error": str(err)}, 404) @jwt_required def post(self, target_os): 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}"' diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index 7bd9a314d..aed4d07f3 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -1 +1,3 @@ from .single_file_repository import SingleFileRepository +from .mock_file_repository import MockFileRepository, FILE_CONTENTS, FILE_NAME +from .open_error_file_repository import OpenErrorFileRepository diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py new file mode 100644 index 000000000..782c9838b --- /dev/null +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -0,0 +1,28 @@ +import io +from typing import BinaryIO + +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository + +FILE_NAME = "test_file" +FILE_CONTENTS = b"HelloWorld!" + + +class MockFileRepository(IFileRepository): + def __init__(self): + self._file = io.BytesIO(FILE_CONTENTS) + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + pass + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + if unsafe_file_name != FILE_NAME: + raise repository.FileNotFoundError() + + return self._file + + def delete_file(self, unsafe_file_name: str): + pass + + def delete_all_files(self): + pass diff --git a/monkey/tests/monkey_island/open_error_file_repository.py b/monkey/tests/monkey_island/open_error_file_repository.py new file mode 100644 index 000000000..c13559613 --- /dev/null +++ b/monkey/tests/monkey_island/open_error_file_repository.py @@ -0,0 +1,10 @@ +from typing import BinaryIO + +from monkey_island.cc.repository import RetrievalError + +from . import MockFileRepository + + +class OpenErrorFileRepository(MockFileRepository): + def open_file(self, unsafe_file_name: str) -> BinaryIO: + raise RetrievalError("Error retrieving file") diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index d00ee9ae9..462969acb 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -1,7 +1,8 @@ import io from typing import BinaryIO -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository class SingleFileRepository(IFileRepository): @@ -13,7 +14,7 @@ class SingleFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: - raise FileRetrievalError() + raise repository.FileNotFoundError() return self._file def delete_file(self, unsafe_file_name: str): diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py index 7ad066623..93a80dea6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py @@ -1,8 +1,9 @@ +import pytest from tests.common.example_agent_configuration import AGENT_CONFIGURATION -from tests.monkey_island import SingleFileRepository +from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema -from monkey_island.cc.repository import FileAgentConfigurationRepository +from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError def test_store_agent_config(): @@ -24,3 +25,10 @@ def test_get_default_agent_config(): retrieved_agent_configuration = repository.get_configuration() assert retrieved_agent_configuration == default_agent_configuration + + +def test_get_agent_config_retrieval_error(): + repository = FileAgentConfigurationRepository(OpenErrorFileRepository()) + + with pytest.raises(RetrievalError): + repository.get_configuration() diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index cf506ffd0..64541bcbe 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -1,10 +1,12 @@ import io from pathlib import Path +from unittest.mock import Mock, patch import pytest from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions -from monkey_island.cc.repository import FileRetrievalError, LocalStorageFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import LocalStorageFileRepository from monkey_island.cc.server_utils.file_utils import is_windows_os @@ -131,5 +133,13 @@ def test_remove_nonexistant_file(tmp_path): def test_open_nonexistant_file(tmp_path): fss = LocalStorageFileRepository(tmp_path) - with pytest.raises(FileRetrievalError): + with pytest.raises(repository.FileNotFoundError): fss.open_file("nonexistant_file.txt") + + +def test_open_locked_file(tmp_path, monkeypatch): + fss = LocalStorageFileRepository(tmp_path) + + with patch("builtins.open", Mock(side_effect=Exception())): + with pytest.raises(repository.RetrievalError): + fss.open_file("locked_file.txt") diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/__init__.py b/monkey/tests/unit_tests/monkey_island/cc/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index a40766d5e..26fe24821 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -2,11 +2,14 @@ from unittest.mock import MagicMock import flask_jwt_extended import pytest +from tests.common import StubDIContainer +from tests.monkey_island import OpenErrorFileRepository from tests.unit_tests.monkey_island.conftest import init_mock_app import monkey_island.cc.app import monkey_island.cc.resources.auth.auth import monkey_island.cc.resources.island_mode +from monkey_island.cc.repository import IFileRepository @pytest.fixture @@ -38,3 +41,12 @@ def get_mock_app(container): flask_jwt_extended.JWTManager(app) return app + + +@pytest.fixture +def open_error_flask_client(build_flask_client): + container = StubDIContainer() + container.register(IFileRepository, OpenErrorFileRepository) + + with build_flask_client(container) as flask_client: + yield flask_client diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index cf2a109b9..31ae0309a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -1,36 +1,11 @@ -import io -from typing import BinaryIO - import pytest from tests.common import StubDIContainer +from tests.monkey_island import FILE_CONTENTS, FILE_NAME, MockFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload -FILE_NAME = "test_file" -FILE_CONTENTS = b"HelloWorld!" - - -class MockFileRepository(IFileRepository): - def __init__(self): - self._file = io.BytesIO(FILE_CONTENTS) - - def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): - pass - - def open_file(self, unsafe_file_name: str) -> BinaryIO: - if unsafe_file_name != FILE_NAME: - raise FileRetrievalError() - - return self._file - - def delete_file(self, unsafe_file_name: str): - pass - - def delete_all_files(self): - pass - @pytest.fixture def flask_client(build_flask_client): @@ -56,3 +31,11 @@ def test_file_download_endpoint_404(tmp_path, flask_client): resp = flask_client.get(download_url) assert resp.status_code == 404 + + +def test_file_download_endpoint_500(tmp_path, open_error_flask_client): + download_url = get_url_for_resource(PBAFileDownload, filename="test") + + resp = open_error_flask_client.get(download_url) + + assert resp.status_code == 500 diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index 50da88915..e4cb5a487 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -98,6 +98,15 @@ def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_conf assert resp.status_code == 404 +@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) +def test_file_download_endpoint_500(open_error_flask_client, pba_os): + url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") + + resp = open_error_flask_client.get(url) + + assert resp.status_code == 500 + + @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_endpoint( flask_client, pba_os, mock_get_config_value, mock_set_config_value