From 3cb678ad3274620de549ec16f2b8759f59a8291c Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 20 Jun 2022 14:55:26 +0200 Subject: [PATCH 01/15] Island: Raise distinct errors when openning a file IFileRepository now distincts between file not found and a file that could not be retrieved --- .../repository/file_storage/i_file_repository.py | 3 ++- .../file_storage/local_storage_file_repository.py | 12 +++++++----- .../test_local_storage_file_repository.py | 14 +++++++++++++- 3 files changed, 22 insertions(+), 7 deletions(-) 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..5069b36d7 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 @@ -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 FileRetrievalError: if the file cannot be retrieved + :raises FileNotFoundError: if the file cannot be found """ 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..a23b25228 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 @@ -1,3 +1,4 @@ +import errno import logging import shutil from pathlib import Path @@ -44,11 +45,12 @@ class LocalStorageFileRepository(IFileRepository): 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 + if err.errno == errno.ENOENT: + raise FileNotFoundError(f"No file {safe_file_path} found: {err}") from err + else: + raise FileRetrievalError( + f"Failed to retrieve file {safe_file_path}: {err}" + ) from err def delete_file(self, unsafe_file_name: str): safe_file_path = self._get_safe_file_path(unsafe_file_name) 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..804b03c79 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,5 +1,7 @@ +import errno 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 @@ -131,5 +133,15 @@ 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(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=OSError(errno.EIO, "File could not be retrieved")) + ): + with pytest.raises(FileRetrievalError): + fss.open_file("locked_file.txt") From c424262f12b841175f68a9d64f33e2a3e629cc9f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 08:22:47 -0400 Subject: [PATCH 02/15] Island: Improve description of return value for get_configuration() --- .../cc/repository/i_agent_configuration_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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..de7aa4ef2 100644 --- a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py @@ -14,7 +14,7 @@ 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. + configuration if the repository is empty :raises RetrievalError: if the configuration can not be retrieved """ pass From cd34cd5eae4d2fb22958ff20a1172ba6d55b5e87 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 08:49:36 -0400 Subject: [PATCH 03/15] Island: Add repository.FileNotFoundError --- .../monkey_island/cc/repository/__init__.py | 2 +- .../cc/repository/agent_binary_repository.py | 7 +++++-- .../cc/repository/file_storage/__init__.py | 2 +- .../file_storage/i_file_repository.py | 4 ++-- .../local_storage_file_repository.py | 20 ++++++++++--------- .../cc/resources/pba_file_download.py | 13 +++++++----- .../cc/resources/pba_file_upload.py | 13 +++++++----- .../monkey_island/single_file_repository.py | 5 +++-- .../test_local_storage_file_repository.py | 12 +++++------ .../cc/resources/test_pba_file_download.py | 4 ++-- 10 files changed, 46 insertions(+), 36 deletions(-) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index cc89c5c57..a44e7b1ad 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,5 +1,5 @@ from .errors import RetrievalError -from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository +from .file_storage import FileNotFoundError, IFileRepository, LocalStorageFileRepository from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError from .agent_binary_repository import AgentBinaryRepository from .i_agent_configuration_repository import IAgentConfigurationRepository diff --git a/monkey/monkey_island/cc/repository/agent_binary_repository.py b/monkey/monkey_island/cc/repository/agent_binary_repository.py index 4c1334f71..a09d88868 100644 --- a/monkey/monkey_island/cc/repository/agent_binary_repository.py +++ b/monkey/monkey_island/cc/repository/agent_binary_repository.py @@ -1,6 +1,8 @@ from typing import BinaryIO -from . import AgentRetrievalError, FileRetrievalError, IAgentBinaryRepository, IFileRepository +from monkey_island.cc import repository + +from . import AgentRetrievalError, IAgentBinaryRepository, IFileRepository LINUX_AGENT_FILE_NAME = "monkey-linux-64" WINDOWS_AGENT_FILE_NAME = "monkey-windows-64.exe" @@ -20,7 +22,8 @@ class AgentBinaryRepository(IAgentBinaryRepository): try: agent_binary = self._file_repository.open_file(filename) return agent_binary - except FileRetrievalError as err: + # TODO: Reevaluate this + except repository.FileNotFoundError as err: raise AgentRetrievalError( 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_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 5069b36d7..681b3d7cf 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,8 +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 retrieved :raises FileNotFoundError: if the file cannot be found + :raises RetrievalError: if the file 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 a23b25228..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 @@ -1,13 +1,13 @@ -import errno import logging import shutil 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__) @@ -44,13 +44,15 @@ class LocalStorageFileRepository(IFileRepository): try: logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") - except OSError as err: - if err.errno == errno.ENOENT: - raise FileNotFoundError(f"No file {safe_file_path} found: {err}") from err - else: - 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/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index 85395d5e3..3b7e92047 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,9 @@ 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: + # TODO: Do we need to log? Or will flask handle it when we `make_response()`? + logger.error(str(err)) + return make_response({"error": str(err)}, 404) + + # TODO: Add unit tests that test 404 vs 500 errors diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 9d5cce8b1..833086998 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,12 @@ 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: + # TODO: Do we need to log? Or will flask handle it when we `make_response()`? + logger.error(str(err)) + return make_response({"error": str(err)}, 404) + + # TODO: Add unit tests that test 404 vs 500 errors @jwt_required def post(self, target_os): 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_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index 804b03c79..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,4 +1,3 @@ -import errno import io from pathlib import Path from unittest.mock import Mock, patch @@ -6,7 +5,8 @@ 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 @@ -133,15 +133,13 @@ def test_remove_nonexistant_file(tmp_path): def test_open_nonexistant_file(tmp_path): fss = LocalStorageFileRepository(tmp_path) - with pytest.raises(FileNotFoundError): + 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=OSError(errno.EIO, "File could not be retrieved")) - ): - with pytest.raises(FileRetrievalError): + 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/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index cf2a109b9..caf86edab 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 @@ -5,7 +5,7 @@ import pytest from tests.common import StubDIContainer 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 FileNotFoundError, IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload FILE_NAME = "test_file" @@ -21,7 +21,7 @@ class MockFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: if unsafe_file_name != FILE_NAME: - raise FileRetrievalError() + raise FileNotFoundError() return self._file From 44795531b850e419f1b6496bc1fca03aa945688e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 21 Jun 2022 15:04:34 +0200 Subject: [PATCH 04/15] Island: Remove logging TODOs for pba_file upload/download Resources should log the errors --- monkey/monkey_island/cc/resources/pba_file_download.py | 1 - monkey/monkey_island/cc/resources/pba_file_upload.py | 1 - 2 files changed, 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index 3b7e92047..39acb5d2c 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -26,7 +26,6 @@ class PBAFileDownload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") except repository.FileNotFoundError as err: - # TODO: Do we need to log? Or will flask handle it when we `make_response()`? 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 833086998..5f46273c4 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -59,7 +59,6 @@ class FileUpload(AbstractResource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") except repository.FileNotFoundError as err: - # TODO: Do we need to log? Or will flask handle it when we `make_response()`? logger.error(str(err)) return make_response({"error": str(err)}, 404) From bcc5265a998da21f09967b4736c4f7bdd63fdf13 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 08:58:21 -0400 Subject: [PATCH 05/15] UT: Add test_file_download_endpoint_500() for PBAFileDownload --- .../cc/resources/test_pba_file_download.py | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) 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 caf86edab..067dcd102 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 @@ -5,7 +5,8 @@ import pytest from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource -from monkey_island.cc.repository import FileNotFoundError, IFileRepository +from monkey_island.cc import repository +from monkey_island.cc.repository import IFileRepository, RetrievalError from monkey_island.cc.resources.pba_file_download import PBAFileDownload FILE_NAME = "test_file" @@ -21,7 +22,7 @@ class MockFileRepository(IFileRepository): def open_file(self, unsafe_file_name: str) -> BinaryIO: if unsafe_file_name != FILE_NAME: - raise FileNotFoundError() + raise repository.FileNotFoundError() return self._file @@ -56,3 +57,25 @@ def test_file_download_endpoint_404(tmp_path, flask_client): resp = flask_client.get(download_url) assert resp.status_code == 404 + + +class OpenErrorFileRepository(MockFileRepository): + def open_file(self, unsafe_file_name: str) -> BinaryIO: + raise RetrievalError("Error retrieving file") + + +@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 + + +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 From bf2f58aacedd3e8895f647eae9f3d89b56b84fef Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 08:59:31 -0400 Subject: [PATCH 06/15] UT: Add __init__.py --- monkey/tests/unit_tests/monkey_island/cc/resources/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/__init__.py 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 From f973c9d6e970d16396eae464da36780d46d290c7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:05:42 -0400 Subject: [PATCH 07/15] UT: Extract MockFileRepository into its own module --- .../cc/resources/mock_file_repository.py | 28 +++++++++++++++++++ .../cc/resources/test_pba_file_download.py | 25 +---------------- 2 files changed, 29 insertions(+), 24 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py new file mode 100644 index 000000000..782c9838b --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/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/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 067dcd102..885f69609 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,13 @@ -import io from typing import BinaryIO import pytest from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource -from monkey_island.cc import repository from monkey_island.cc.repository import IFileRepository, RetrievalError 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 repository.FileNotFoundError() - - return self._file - - def delete_file(self, unsafe_file_name: str): - pass - - def delete_all_files(self): - pass +from .mock_file_repository import FILE_CONTENTS, FILE_NAME, MockFileRepository @pytest.fixture From 8939ca21066cf429808e8cbf3432eb684fe5475f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:08:32 -0400 Subject: [PATCH 08/15] UT: Extract open_error_flask_client into conftest.py --- .../monkey_island/cc/resources/conftest.py | 19 +++++++++++++++++++ .../cc/resources/test_pba_file_download.py | 18 +----------------- 2 files changed, 20 insertions(+), 17 deletions(-) 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..532282e8b 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,12 +1,17 @@ +from typing import BinaryIO from unittest.mock import MagicMock import flask_jwt_extended import pytest +from tests.common import StubDIContainer 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, RetrievalError + +from .mock_file_repository import MockFileRepository @pytest.fixture @@ -38,3 +43,17 @@ def get_mock_app(container): flask_jwt_extended.JWTManager(app) return app + + +class OpenErrorFileRepository(MockFileRepository): + def open_file(self, unsafe_file_name: str) -> BinaryIO: + raise RetrievalError("Error retrieving file") + + +@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 885f69609..eb189b4e7 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,10 +1,8 @@ -from typing import BinaryIO - import pytest from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource -from monkey_island.cc.repository import IFileRepository, RetrievalError +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload from .mock_file_repository import FILE_CONTENTS, FILE_NAME, MockFileRepository @@ -36,20 +34,6 @@ def test_file_download_endpoint_404(tmp_path, flask_client): assert resp.status_code == 404 -class OpenErrorFileRepository(MockFileRepository): - def open_file(self, unsafe_file_name: str) -> BinaryIO: - raise RetrievalError("Error retrieving file") - - -@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 - - def test_file_download_endpoint_500(tmp_path, open_error_flask_client): download_url = get_url_for_resource(PBAFileDownload, filename="test") From 63a2527f3f233ac7aff3c5c1b317b2c916270cde Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:10:40 -0400 Subject: [PATCH 09/15] UT: Add test_file_download_endpoint_500() for PBAFileUpload --- .../monkey_island/cc/resources/test_pba_file_upload.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 From 22b22c5f0a7ce5a24d45bb1df00b62808e450670 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:11:54 -0400 Subject: [PATCH 10/15] Island: Remove testing TODOs from PBAFile{Download,Upload} --- monkey/monkey_island/cc/resources/pba_file_download.py | 2 -- monkey/monkey_island/cc/resources/pba_file_upload.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index 39acb5d2c..d8359049a 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -28,5 +28,3 @@ class PBAFileDownload(AbstractResource): except repository.FileNotFoundError as err: logger.error(str(err)) return make_response({"error": str(err)}, 404) - - # TODO: Add unit tests that test 404 vs 500 errors diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 5f46273c4..277804a72 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -62,8 +62,6 @@ class FileUpload(AbstractResource): logger.error(str(err)) return make_response({"error": str(err)}, 404) - # TODO: Add unit tests that test 404 vs 500 errors - @jwt_required def post(self, target_os): """ From 47df2575459f2e9a1392d6f8d0f44dbf97ae1ee4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:18:42 -0400 Subject: [PATCH 11/15] 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}"' From 4de9f3cb6d2181abe3f32ed6bdb04769be1f3829 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:35:23 -0400 Subject: [PATCH 12/15] UT: Move OpenErrorFileRepository into tests/monkey_island --- monkey/tests/monkey_island/__init__.py | 2 ++ .../mock_file_repository.py | 0 .../tests/monkey_island/open_error_file_repository.py | 10 ++++++++++ .../unit_tests/monkey_island/cc/resources/conftest.py | 11 ++--------- .../cc/resources/test_pba_file_download.py | 3 +-- 5 files changed, 15 insertions(+), 11 deletions(-) rename monkey/tests/{unit_tests/monkey_island/cc/resources => monkey_island}/mock_file_repository.py (100%) create mode 100644 monkey/tests/monkey_island/open_error_file_repository.py 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/unit_tests/monkey_island/cc/resources/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py similarity index 100% rename from monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py rename to monkey/tests/monkey_island/mock_file_repository.py 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/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index 532282e8b..26fe24821 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,17 +1,15 @@ -from typing import BinaryIO 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, RetrievalError - -from .mock_file_repository import MockFileRepository +from monkey_island.cc.repository import IFileRepository @pytest.fixture @@ -45,11 +43,6 @@ def get_mock_app(container): return app -class OpenErrorFileRepository(MockFileRepository): - def open_file(self, unsafe_file_name: str) -> BinaryIO: - raise RetrievalError("Error retrieving file") - - @pytest.fixture def open_error_flask_client(build_flask_client): container = StubDIContainer() 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 eb189b4e7..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,12 +1,11 @@ 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 IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload -from .mock_file_repository import FILE_CONTENTS, FILE_NAME, MockFileRepository - @pytest.fixture def flask_client(build_flask_client): From c008db4cf2857779eb0f062748da1e313d3c98d6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 09:40:00 -0400 Subject: [PATCH 13/15] Island: Raise RetrievalError in FileAgentConfigurationRepository --- .../file_agent_configuration_repository.py | 7 ++++--- .../test_file_agent_configuration_repository.py | 12 ++++++++++-- 2 files changed, 14 insertions(+), 5 deletions(-) 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/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() From 3bd977ed551cd5a17cf41ec7f30b53b537e463fb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 10:34:27 -0400 Subject: [PATCH 14/15] Island: Improve IFileRepository.open_file() docstring --- .../cc/repository/file_storage/i_file_repository.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 681b3d7cf..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 @@ -30,8 +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 FileNotFoundError: if the file cannot be found - :raises RetrievalError: if the file cannot be retrieved + :raises FileNotFoundError: if the file does not exist + :raises RetrievalError: if the file exists but cannot be retrieved """ pass From 838a2e742cfb11ace9b832096824e6d8d89466e3 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 10:36:56 -0400 Subject: [PATCH 15/15] Island: Change can -> could in get_configuration() --- .../cc/repository/i_agent_configuration_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 de7aa4ef2..1e85fadf5 100644 --- a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py @@ -15,7 +15,7 @@ class IAgentConfigurationRepository(ABC): :return: The agent configuration as retrieved from the repository, or the default configuration if the repository is empty - :raises RetrievalError: if the configuration can not be retrieved + :raises RetrievalError: if the configuration could not be retrieved """ pass