From 3c1e25b88cc96c5b8dc2f5349ed1f2bb41f3cb2b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 22 Apr 2022 09:35:41 -0400 Subject: [PATCH 01/35] UT: Move Linux directory permissions check to a utility function --- monkey/tests/monkey_island/utils.py | 12 ++++++++++++ .../monkey_island/cc/server_utils/test_file_utils.py | 8 ++------ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/monkey/tests/monkey_island/utils.py b/monkey/tests/monkey_island/utils.py index dd781f85d..8ab767774 100644 --- a/monkey/tests/monkey_island/utils.py +++ b/monkey/tests/monkey_island/utils.py @@ -7,6 +7,9 @@ if is_windows_os(): FULL_CONTROL = 2032127 ACE_ACCESS_MODE_GRANT_ACCESS = win32security.GRANT_ACCESS ACE_INHERIT_OBJECT_AND_CONTAINER = 3 +else: + import os + import stat def _get_acl_and_sid_from_path(path: str): @@ -33,3 +36,12 @@ def assert_windows_permissions(path: str): assert ace_sid == user_sid assert ace_permissions == FULL_CONTROL and ace_access_mode == ACE_ACCESS_MODE_GRANT_ACCESS assert ace_inheritance == ACE_INHERIT_OBJECT_AND_CONTAINER + + +def assert_linux_permissions(path: str): + st = os.stat(path) + + expected_mode = stat.S_IRWXU + actual_mode = st.st_mode & (stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) + + assert expected_mode == actual_mode diff --git a/monkey/tests/unit_tests/monkey_island/cc/server_utils/test_file_utils.py b/monkey/tests/unit_tests/monkey_island/cc/server_utils/test_file_utils.py index 6605673d0..7d05e29c2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/server_utils/test_file_utils.py +++ b/monkey/tests/unit_tests/monkey_island/cc/server_utils/test_file_utils.py @@ -2,7 +2,7 @@ import os import stat import pytest -from tests.monkey_island.utils import assert_windows_permissions +from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions from monkey_island.cc.server_utils.file_utils import ( create_secure_directory, @@ -39,12 +39,8 @@ def test_create_secure_directory__no_parent_dir(test_path_nested): @pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.") def test_create_secure_directory__perm_linux(test_path): create_secure_directory(test_path) - st = os.stat(test_path) - expected_mode = stat.S_IRWXU - actual_mode = st.st_mode & (stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) - - assert expected_mode == actual_mode + assert_linux_permissions(test_path) @pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.") From 80a305ea81b7e6a505a63422ede9880098da2bf3 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 22 Apr 2022 12:30:13 -0400 Subject: [PATCH 02/35] Island: Add IFileStorageService --- monkey/monkey_island/cc/services/__init__.py | 1 + .../cc/services/i_file_storage_service.py | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 monkey/monkey_island/cc/services/i_file_storage_service.py diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index b3dcba9de..15d0b801d 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,2 +1,3 @@ from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore +from .i_file_storage_service import IFileStorageService diff --git a/monkey/monkey_island/cc/services/i_file_storage_service.py b/monkey/monkey_island/cc/services/i_file_storage_service.py new file mode 100644 index 000000000..a27cbf6ad --- /dev/null +++ b/monkey/monkey_island/cc/services/i_file_storage_service.py @@ -0,0 +1,45 @@ +import abc +from typing import BinaryIO + + +class IFileStorageService(metaclass=abc.ABCMeta): + """ + A service that allows the storage and retrieval of individual files. + """ + + @abc.abstractmethod + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + """ + Save a file, identified by a name + + :param unsafe_file_name: An unsanitized file name that will identify the file + :param file_contents: The data to be stored in the file + """ + pass + + @abc.abstractmethod + def open_file(self, unsafe_file_name: str) -> BinaryIO: + """ + Open a file and return a file-like object + + :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 + :rtype: io.BinaryIO + """ + pass + + @abc.abstractmethod + def delete_file(self, unsafe_file_name: str): + """ + Delete a file + + :param unsafe_file_name: An unsanitized file name that identifies the file to be deleted + """ + pass + + @abc.abstractmethod + def delete_all_files(self): + """ + Delete all files that have been stored using this service. + """ + pass From 88df935c770481ab28162038d0d4788f417508d0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 22 Apr 2022 12:30:32 -0400 Subject: [PATCH 03/35] Island: Add DirectoryFileStorageService --- monkey/monkey_island/cc/services/__init__.py | 1 + .../directory_file_storage_service.py | 54 +++++++++ .../test_directory_file_storage_service.py | 108 ++++++++++++++++++ 3 files changed, 163 insertions(+) create mode 100644 monkey/monkey_island/cc/services/directory_file_storage_service.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 15d0b801d..2d13fae99 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,3 +1,4 @@ from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore from .i_file_storage_service import IFileStorageService +from .directory_file_storage_service import DirectoryFileStorageService diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py new file mode 100644 index 000000000..30aa2a4d0 --- /dev/null +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -0,0 +1,54 @@ +import shutil +from pathlib import Path +from typing import BinaryIO + +from monkey_island.cc.server_utils.file_utils import create_secure_directory + +from . import IFileStorageService + + +class DirectoryFileStorageService(IFileStorageService): + """ + A implementation of IFileStorageService that reads and writes files from/to the local + filesystem. + """ + + def __init__(self, storage_directory: Path): + """ + :param storage_directory: A Path object representing the directory where files will be + stored. If the directory does not exist, it will be created. + """ + if storage_directory.exists() and not storage_directory.is_dir(): + raise ValueError(f"The provided path must point to a directory: {storage_directory}") + + if not storage_directory.exists(): + create_secure_directory(storage_directory) + + self._storage_directory = storage_directory + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + safe_file_path = self._get_safe_file_path(unsafe_file_name) + + with open(safe_file_path, "wb") as dest: + shutil.copyfileobj(file_contents, dest) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + safe_file_path = self._get_safe_file_path(unsafe_file_name) + return open(safe_file_path, "rb") + + def delete_file(self, unsafe_file_name: str): + safe_file_path = self._get_safe_file_path(unsafe_file_name) + + safe_file_path.unlink() + + def _get_safe_file_path(self, unsafe_file_name: str): + # Remove any path information from the file name. + safe_file_name = Path(unsafe_file_name).resolve().name + + # TODO: Add super paranoid check + + return self._storage_directory / safe_file_name + + def delete_all_files(self): + for file in self._storage_directory.iterdir(): + file.unlink() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py new file mode 100644 index 000000000..7d65489d1 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py @@ -0,0 +1,108 @@ +import io +from pathlib import Path + +import pytest +from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions + +from monkey_island.cc.server_utils.file_utils import is_windows_os +from monkey_island.cc.services import DirectoryFileStorageService + + +def test_error_if_file(tmp_path): + new_file = tmp_path / "new_file.txt" + new_file.write_text("HelloWorld!") + + with pytest.raises(ValueError): + DirectoryFileStorageService(new_file) + + +def test_directory_created(tmp_path): + new_dir = tmp_path / "new_dir" + + DirectoryFileStorageService(new_dir) + + assert new_dir.exists() and new_dir.is_dir() + + +@pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.") +def test_directory_permissions__linux(tmp_path): + new_dir = tmp_path / "new_dir" + + DirectoryFileStorageService(new_dir) + + assert_linux_permissions(new_dir) + + +@pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.") +def test_directory_permissions__windows(tmp_path): + new_dir = tmp_path / "new_dir" + + DirectoryFileStorageService(new_dir) + + assert_windows_permissions(new_dir) + + +def save_file(tmp_path, file_path_prefix=""): + file_name = "test.txt" + file_contents = "Hello World!" + expected_file_path = tmp_path / file_name + + fss = DirectoryFileStorageService(tmp_path) + fss.save_file(Path(file_path_prefix) / file_name, io.BytesIO(file_contents.encode())) + + assert expected_file_path.is_file() + assert expected_file_path.read_text() == file_contents + + +def delete_file(tmp_path, file_path_prefix=""): + file_name = "file.txt" + file = tmp_path / file_name + file.touch() + assert file.is_file() + + fss = DirectoryFileStorageService(tmp_path) + fss.delete_file(Path(file_path_prefix) / file_name) + + assert not file.exists() + + +def open_file(tmp_path, file_path_prefix=""): + file_name = "test.txt" + expected_file_contents = "Hello World!" + expected_file_path = tmp_path / file_name + expected_file_path.write_text(expected_file_contents) + + fss = DirectoryFileStorageService(tmp_path) + with fss.open_file(Path(file_path_prefix) / file_name) as f: + actual_file_contents = f.read() + + assert actual_file_contents == expected_file_contents.encode() + + +@pytest.mark.parametrize("fn", [save_file, open_file, delete_file]) +def test_fn(tmp_path, fn): + fn(tmp_path) + + +@pytest.mark.parametrize("fn", [save_file, open_file, delete_file]) +def test_fn__ignore_relative_path(tmp_path, fn): + fn(tmp_path, "../../") + + +@pytest.mark.parametrize("fn", [save_file, open_file, delete_file]) +def test_fn__ignore_absolute_path(tmp_path, fn): + if is_windows_os(): + fn(tmp_path, "C:\\Windows") + else: + fn(tmp_path, "/home/") + + +def test_remove_all_files(tmp_path): + for filename in ["1.txt", "2.txt", "3.txt"]: + (tmp_path / filename).touch() + + fss = DirectoryFileStorageService(tmp_path) + fss.delete_all_files() + + for file in tmp_path.iterdir(): + assert False, f"{tmp_path} was expected to be empty, but contained files" From 3355455c9bb59ad0477aacb56700c08bc9141405 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 10:21:44 -0400 Subject: [PATCH 04/35] Project: Ignore .mypy_cache directories --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index d5cbe3b30..a0f70a939 100644 --- a/.gitignore +++ b/.gitignore @@ -101,3 +101,6 @@ venv/ # Hugo .hugo_build.lock + +# mypy +.mypy_cache From 6cfdcaaec707a2ba20ef3650cbd48917ceb0f8e1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 10:23:37 -0400 Subject: [PATCH 05/35] UT: Change flask_client() fixture from session- to function-scoped --- monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 eeef5b383..4779bd0e2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -9,7 +9,7 @@ import monkey_island.cc.resources.island_mode from monkey_island.cc.services.representations import output_json -@pytest.fixture(scope="session") +@pytest.fixture def flask_client(monkeypatch_session): monkeypatch_session.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) From c03a5aac4bb014002be179caabc3ebc630d24c67 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 10:26:32 -0400 Subject: [PATCH 06/35] Island: Pass the data directory to init_app_resources() --- monkey/monkey_island/cc/app.py | 7 ++++--- monkey/monkey_island/cc/server_setup.py | 2 +- .../unit_tests/monkey_island/cc/resources/conftest.py | 8 ++++---- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 78f09b4de..77d52ac8c 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,6 +1,7 @@ import os import uuid from datetime import timedelta +from pathlib import Path import flask_restful from flask import Flask, Response, send_from_directory @@ -111,7 +112,7 @@ def init_app_url_rules(app): app.add_url_rule("/", "serve_static_file", serve_static_file) -def init_api_resources(api): +def init_api_resources(api, data_dir: Path): api.add_resource(Root, "/api") api.add_resource(Registration, "/api/registration") api.add_resource(Authenticate, "/api/auth") @@ -168,7 +169,7 @@ def init_api_resources(api): api.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") -def init_app(mongo_url): +def init_app(mongo_url, data_dir: Path): app = Flask(__name__) api = flask_restful.Api(app) @@ -177,6 +178,6 @@ def init_app(mongo_url): init_app_config(app, mongo_url) init_app_services(app) init_app_url_rules(app) - init_api_resources(api) + init_api_resources(api, data_dir) return app diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 98f29de10..072cb32a8 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -129,7 +129,7 @@ def _configure_gevent_exception_handling(data_dir): def _start_island_server(should_setup_only, config_options: IslandConfigOptions): populate_exporter_list() - app = init_app(mongo_setup.MONGO_URL) + app = init_app(mongo_setup.MONGO_URL, Path(config_options.data_dir)) if should_setup_only: logger.warning("Setup only flag passed. Exiting.") 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 4779bd0e2..903528be2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -10,14 +10,14 @@ from monkey_island.cc.services.representations import output_json @pytest.fixture -def flask_client(monkeypatch_session): +def flask_client(monkeypatch_session, tmp_path): monkeypatch_session.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) - with mock_init_app().test_client() as client: + with mock_init_app(tmp_path).test_client() as client: yield client -def mock_init_app(): +def mock_init_app(data_dir): app = Flask(__name__) app.config["SECRET_KEY"] = "test_key" @@ -25,7 +25,7 @@ def mock_init_app(): api.representations = {"application/json": output_json} monkey_island.cc.app.init_app_url_rules(app) - monkey_island.cc.app.init_api_resources(api) + monkey_island.cc.app.init_api_resources(api, data_dir) flask_jwt_extended.JWTManager(app) From d1e18e9dbda96d681e1e29a7707d01bc60dffbbf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 10:57:17 -0400 Subject: [PATCH 07/35] Island: Use IFileStorageService in PBAFileDownload resource --- monkey/monkey_island/cc/app.py | 7 ++++- .../cc/resources/pba_file_download.py | 26 ++++++++++++++----- .../cc/resources/test_pba_file_download.py | 17 ++++++++++++ 3 files changed, 43 insertions(+), 7 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 77d52ac8c..17e86fc23 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -47,6 +47,7 @@ from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFinding from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.custom_json_encoder import CustomJSONEncoder +from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService from monkey_island.cc.services.representations import output_json @@ -149,7 +150,11 @@ def init_api_resources(api, data_dir: Path): api.add_resource(TelemetryFeed, "/api/telemetry-feed") api.add_resource(Log, "/api/log") api.add_resource(IslandLog, "/api/log/island/download") - api.add_resource(PBAFileDownload, "/api/pba/download/") + api.add_resource( + PBAFileDownload, + "/api/pba/download/", + resource_class_kwargs={"file_storage_service": DirectoryFileStorageService(data_dir)}, + ) api.add_resource( FileUpload, "/api/file-upload/", diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index df9766ed6..ba5857c53 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -1,7 +1,11 @@ -import flask_restful -from flask import send_from_directory +import logging -from monkey_island.cc.services.post_breach_files import PostBreachFilesService +import flask_restful +from flask import make_response, send_file + +from monkey_island.cc.services import IFileStorageService + +logger = logging.getLogger(__file__) class PBAFileDownload(flask_restful.Resource): @@ -9,7 +13,17 @@ class PBAFileDownload(flask_restful.Resource): File download endpoint used by monkey to download user's PBA file """ + def __init__(self, file_storage_service: IFileStorageService): + self._file_storage_service = file_storage_service + # Used by monkey. can't secure. - def get(self, filename): - custom_pba_dir = PostBreachFilesService.get_custom_pba_directory() - return send_from_directory(custom_pba_dir, filename) + def get(self, filename: str): + try: + file = self._file_storage_service.open_file(filename) + + # `send_file()` handles the closing of the open file. + return send_file(file, mimetype="application/octet-stream") + except OSError as ex: + error_msg = f"Failed to open file {filename}: {ex}" + logger.error(error_msg) + return make_response({"error": error_msg}, 404) 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 new file mode 100644 index 000000000..333ff4cf1 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -0,0 +1,17 @@ +def test_file_download_endpoint(tmp_path, flask_client): + file_contents = "HelloWorld!" + file_name = "test_file" + (tmp_path / file_name).write_text(file_contents) + + resp = flask_client.get(f"/api/pba/download/{file_name}") + + assert resp.status_code == 200 + assert next(resp.response).decode() == file_contents + + +def test_file_download_endpoint_404(tmp_path, flask_client): + file_name = "nonexistant_file" + + resp = flask_client.get(f"/api/pba/download/{file_name}") + + assert resp.status_code == 404 From d157bf7a40b1e4babc3783da64752390dd0babaf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 10:59:32 -0400 Subject: [PATCH 08/35] Island: Use IFileStorageService in FileUpload resource --- monkey/monkey_island/cc/app.py | 7 ++- .../cc/resources/pba_file_upload.py | 46 ++++++++++++++----- .../cc/resources/test_pba_file_download.py | 2 +- 3 files changed, 41 insertions(+), 14 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 17e86fc23..beb9e1468 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -150,17 +150,22 @@ def init_api_resources(api, data_dir: Path): api.add_resource(TelemetryFeed, "/api/telemetry-feed") api.add_resource(Log, "/api/log") api.add_resource(IslandLog, "/api/log/island/download") + + # This is temporary until we get DI all worked out. + file_storage_service = DirectoryFileStorageService(data_dir / "custom_pbas") api.add_resource( PBAFileDownload, "/api/pba/download/", - resource_class_kwargs={"file_storage_service": DirectoryFileStorageService(data_dir)}, + resource_class_kwargs={"file_storage_service": file_storage_service}, ) api.add_resource( FileUpload, "/api/file-upload/", "/api/file-upload/?load=", "/api/file-upload/?restore=", + resource_class_kwargs={"file_storage_service": file_storage_service}, ) + api.add_resource(PropagationCredentials, "/api/propagation-credentials/") api.add_resource(RemoteRun, "/api/remote-monkey") api.add_resource(VersionUpdate, "/api/version-update") diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 68be20e9a..bb2608872 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -1,15 +1,19 @@ import copy +import logging from http import HTTPStatus import flask_restful -from flask import Response, request, send_from_directory +from flask import Response, make_response, request, send_file from werkzeug.datastructures import FileStorage from werkzeug.utils import secure_filename from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.services import IFileStorageService from monkey_island.cc.services.config import ConfigService -from monkey_island.cc.services.post_breach_files import PostBreachFilesService + +logger = logging.getLogger(__file__) + # Front end uses these strings to identify which files to work with (linux or windows) LINUX_PBA_TYPE = "PBAlinux" @@ -21,6 +25,14 @@ class FileUpload(flask_restful.Resource): File upload endpoint used to exchange files with filepond component on the front-end """ + def __init__(self, file_storage_service: IFileStorageService): + self._file_storage_service = file_storage_service + + # TODO: Fix references/coupling to filepond + # TODO: Replace "file_type" with "target_os" or similar + # TODO: Prefix private functions with "_" + # TODO: Add comment explaining why this is basically a duplicate of the endpoint in the + # PBAFileDownload resource. @jwt_required def get(self, file_type): """ @@ -33,10 +45,20 @@ class FileUpload(flask_restful.Resource): # Verify that file_name is indeed a file from config if file_type == LINUX_PBA_TYPE: + # TODO: Make these paths Tuples so we don't need to copy them filename = ConfigService.get_config_value(copy.deepcopy(PBA_LINUX_FILENAME_PATH)) else: filename = ConfigService.get_config_value(copy.deepcopy(PBA_WINDOWS_FILENAME_PATH)) - return send_from_directory(PostBreachFilesService.get_custom_pba_directory(), filename) + + try: + file = self._file_storage_service.open_file(filename) + + # `send_file()` handles the closing of the open file. + return send_file(file, mimetype="application/octet-stream") + except OSError as ex: + error_msg = f"Failed to open file {filename}: {ex}" + logger.error(error_msg) + return make_response({"error": error_msg}, 404) @jwt_required def post(self, file_type): @@ -48,15 +70,17 @@ class FileUpload(flask_restful.Resource): if self.is_pba_file_type_supported(file_type): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") - filename = FileUpload.upload_pba_file( - request.files["filepond"], (file_type == LINUX_PBA_TYPE) + filename = self.upload_pba_file( + # TODO: This "filepond" string can be changed to be more generic in the `react-filepond` + # component. + request.files["filepond"], + (file_type == LINUX_PBA_TYPE), ) response = Response(response=filename, status=200, mimetype="text/plain") return response - @staticmethod - def upload_pba_file(file_storage: FileStorage, is_linux=True): + def upload_pba_file(self, file_storage: FileStorage, is_linux=True): """ Uploads PBA file to island's file system :param request_: Request object containing PBA file @@ -64,9 +88,7 @@ class FileUpload(flask_restful.Resource): :return: filename string """ filename = secure_filename(file_storage.filename) - file_contents = file_storage.read() - - PostBreachFilesService.save_file(filename, file_contents) + self._file_storage_service.save_file(filename, file_storage.stream) ConfigService.set_config_value( (PBA_LINUX_FILENAME_PATH if is_linux else PBA_WINDOWS_FILENAME_PATH), filename @@ -89,10 +111,10 @@ class FileUpload(flask_restful.Resource): ) filename = ConfigService.get_config_value(filename_path) if filename: - PostBreachFilesService.remove_file(filename) + self._file_storage_service.delete_file(filename) ConfigService.set_config_value(filename_path, "") - return {} + return make_response({}, 200) @staticmethod def is_pba_file_type_supported(file_type: str) -> bool: 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 333ff4cf1..5b6dcad3f 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,7 +1,7 @@ def test_file_download_endpoint(tmp_path, flask_client): file_contents = "HelloWorld!" file_name = "test_file" - (tmp_path / file_name).write_text(file_contents) + (tmp_path / "custom_pbas" / file_name).write_text(file_contents) resp = flask_client.get(f"/api/pba/download/{file_name}") From a487aa40581a0c2b6e7717e8a9738df4e61e4528 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 11:52:04 -0400 Subject: [PATCH 09/35] Island: Modify PostBreachFilesService to wrap IFileStorageService --- monkey/monkey_island/cc/services/__init__.py | 5 +- .../monkey_island/cc/services/initialize.py | 5 +- .../cc/services/post_breach_files.py | 44 ++++--------- .../cc/resources/test_pba_file_upload.py | 6 -- .../cc/services/test_post_breach_files.py | 66 +++++-------------- 5 files changed, 35 insertions(+), 91 deletions(-) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 2d13fae99..3f57b4fc0 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,4 +1,5 @@ -from .authentication.authentication_service import AuthenticationService -from .authentication.json_file_user_datastore import JsonFileUserDatastore from .i_file_storage_service import IFileStorageService from .directory_file_storage_service import DirectoryFileStorageService + +from .authentication.authentication_service import AuthenticationService +from .authentication.json_file_user_datastore import JsonFileUserDatastore diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 64797eff0..f3ff10395 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,3 +1,4 @@ +from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -5,6 +6,8 @@ from . import AuthenticationService, JsonFileUserDatastore def initialize_services(data_dir): - PostBreachFilesService.initialize(data_dir) + # This is temporary until we get DI all worked out. + PostBreachFilesService.initialize(DirectoryFileStorageService(data_dir / "custom_pbas")) + LocalMonkeyRunService.initialize(data_dir) AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) diff --git a/monkey/monkey_island/cc/services/post_breach_files.py b/monkey/monkey_island/cc/services/post_breach_files.py index 8268265a9..9b24ce391 100644 --- a/monkey/monkey_island/cc/services/post_breach_files.py +++ b/monkey/monkey_island/cc/services/post_breach_files.py @@ -1,46 +1,24 @@ import logging -import os -from monkey_island.cc.server_utils.file_utils import create_secure_directory +from monkey_island.cc.services import IFileStorageService logger = logging.getLogger(__name__) +# TODO: This service wraps an IFileStorageService for the sole purpose of making the +# `remove_PBA_files()` method available to the ConfigService. This whole service can be +# removed once ConfigService is refactored to be stateful (it already is but everything is +# still statically/globally scoped) and use dependency injection. class PostBreachFilesService: - DATA_DIR = None - CUSTOM_PBA_DIRNAME = "custom_pbas" + _file_storage_service = None # TODO: A number of these services should be instance objects instead of # static/singleton hybrids. At the moment, this requires invasive refactoring that's # not a priority. @classmethod - def initialize(cls, data_dir): - cls.DATA_DIR = data_dir - custom_pba_dir = cls.get_custom_pba_directory() - create_secure_directory(custom_pba_dir) + def initialize(cls, file_storage_service: IFileStorageService): + cls._file_storage_service = file_storage_service - @staticmethod - def save_file(filename: str, file_contents: bytes): - file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), filename) - with open(file_path, "wb") as f: - f.write(file_contents) - - @staticmethod - def remove_PBA_files(): - for f in os.listdir(PostBreachFilesService.get_custom_pba_directory()): - PostBreachFilesService.remove_file(f) - - @staticmethod - def remove_file(file_name): - file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), file_name) - try: - if os.path.exists(file_path): - os.remove(file_path) - except OSError as e: - logger.error("Can't remove previously uploaded post breach files: %s" % e) - - @staticmethod - def get_custom_pba_directory(): - return os.path.join( - PostBreachFilesService.DATA_DIR, PostBreachFilesService.CUSTOM_PBA_DIRNAME - ) + @classmethod + def remove_PBA_files(cls): + cls._file_storage_service.delete_all_files() 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 a0b8352e2..ca5f0f175 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 @@ -2,7 +2,6 @@ import pytest from tests.utils import raise_ from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE -from monkey_island.cc.services.post_breach_files import PostBreachFilesService TEST_FILE = b"""-----------------------------1 Content-Disposition: form-data; name="filepond" @@ -16,11 +15,6 @@ m0nk3y -----------------------------1--""" -@pytest.fixture(autouse=True) -def custom_pba_directory(tmpdir): - PostBreachFilesService.initialize(tmpdir) - - @pytest.fixture def mock_set_config_value(monkeypatch): monkeypatch.setattr( diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py index 90a649a39..23f42cf79 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py @@ -1,29 +1,31 @@ +import io import os import pytest -from tests.monkey_island.utils import assert_windows_permissions from tests.utils import raise_ -from monkey_island.cc.server_utils.file_utils import is_windows_os +from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService +@pytest.fixture +def file_storage_service(tmp_path): + return DirectoryFileStorageService(tmp_path) + + @pytest.fixture(autouse=True) -def custom_pba_directory(tmpdir): - PostBreachFilesService.initialize(tmpdir) +def post_breach_files_service(file_storage_service): + PostBreachFilesService.initialize(file_storage_service) -def create_custom_pba_file(filename): - PostBreachFilesService.save_file(filename, b"") +def test_remove_pba_files(file_storage_service, tmp_path): + file_storage_service.save_file("linux_file", io.BytesIO(b"")) + file_storage_service.save_file("windows_file", io.BytesIO(b"")) + assert not dir_is_empty(tmp_path) - -def test_remove_pba_files(): - create_custom_pba_file("linux_file") - create_custom_pba_file("windows_file") - - assert not dir_is_empty(PostBreachFilesService.get_custom_pba_directory()) PostBreachFilesService.remove_PBA_files() - assert dir_is_empty(PostBreachFilesService.get_custom_pba_directory()) + + assert dir_is_empty(tmp_path) def dir_is_empty(dir_path): @@ -31,45 +33,11 @@ def dir_is_empty(dir_path): return len(dir_contents) == 0 -@pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.") -def test_custom_pba_dir_permissions_linux(): - st = os.stat(PostBreachFilesService.get_custom_pba_directory()) - - assert st.st_mode == 0o40700 - - -@pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.") -def test_custom_pba_dir_permissions_windows(): - pba_dir = PostBreachFilesService.get_custom_pba_directory() - - assert_windows_permissions(pba_dir) - - -def test_remove_failure(monkeypatch): +def test_remove_failure(file_storage_service, monkeypatch): monkeypatch.setattr(os, "remove", lambda x: raise_(OSError("Permission denied"))) try: - create_custom_pba_file("windows_file") + file_storage_service.save_file("windows_file", io.BytesIO(b"")) PostBreachFilesService.remove_PBA_files() except Exception as ex: pytest.fail(f"Unxepected exception: {ex}") - - -def test_remove_nonexistant_file(monkeypatch): - monkeypatch.setattr(os, "remove", lambda x: raise_(FileNotFoundError("FileNotFound"))) - - try: - PostBreachFilesService.remove_file("/nonexistant/file") - except Exception as ex: - pytest.fail(f"Unxepected exception: {ex}") - - -def test_save_file(): - FILE_NAME = "test_file" - FILE_CONTENTS = b"hello" - PostBreachFilesService.save_file(FILE_NAME, FILE_CONTENTS) - - expected_file_path = os.path.join(PostBreachFilesService.get_custom_pba_directory(), FILE_NAME) - - assert os.path.isfile(expected_file_path) - assert FILE_CONTENTS == open(expected_file_path, "rb").read() From 2f4ffad3f655214ddd9b2dbc3dc7700eebcb9cf8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 12:15:57 -0400 Subject: [PATCH 10/35] Island: Skip directories in DirectoryFileStorageService.delete_all_files --- .../cc/services/directory_file_storage_service.py | 2 +- .../services/test_directory_file_storage_service.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index 30aa2a4d0..48fb95bfc 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -50,5 +50,5 @@ class DirectoryFileStorageService(IFileStorageService): return self._storage_directory / safe_file_name def delete_all_files(self): - for file in self._storage_directory.iterdir(): + for file in filter(lambda f: f.is_file(), self._storage_directory.iterdir()): file.unlink() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py index 7d65489d1..09ed308c5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py @@ -106,3 +106,16 @@ def test_remove_all_files(tmp_path): for file in tmp_path.iterdir(): assert False, f"{tmp_path} was expected to be empty, but contained files" + + +def test_remove_all_files__skip_directories(tmp_path): + test_dir = tmp_path / "test_dir" + test_dir.mkdir() + for filename in ["1.txt", "2.txt", "3.txt"]: + (tmp_path / filename).touch() + + fss = DirectoryFileStorageService(tmp_path) + fss.delete_all_files() + + for file in tmp_path.iterdir(): + assert file.name == test_dir.name From cd8fa699b0ff5f58e02492cb899b5287e7ca0700 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 12:32:59 -0400 Subject: [PATCH 11/35] Common: Move get_all_regular_files_in_directory() to utils.file_utils --- monkey/common/utils/file_utils.py | 5 ++ .../payload/ransomware/file_selectors.py | 2 +- monkey/infection_monkey/utils/dir_utils.py | 4 -- .../directory_file_storage_service.py | 3 +- .../common/utils/test_common_file_utils.py | 37 +++++++++++- .../infection_monkey/utils/test_dir_utils.py | 60 +++---------------- monkey/tests/utils.py | 20 +++++++ 7 files changed, 72 insertions(+), 59 deletions(-) diff --git a/monkey/common/utils/file_utils.py b/monkey/common/utils/file_utils.py index fd2c85ec1..c9cb78139 100644 --- a/monkey/common/utils/file_utils.py +++ b/monkey/common/utils/file_utils.py @@ -1,6 +1,7 @@ import hashlib import os from pathlib import Path +from typing import Iterable class InvalidPath(Exception): @@ -21,3 +22,7 @@ def get_file_sha256_hash(filepath: Path): sha256.update(block) return sha256.hexdigest() + + +def get_all_regular_files_in_directory(dir_path: Path) -> Iterable[Path]: + return filter(lambda f: f.is_file(), dir_path.iterdir()) diff --git a/monkey/infection_monkey/payload/ransomware/file_selectors.py b/monkey/infection_monkey/payload/ransomware/file_selectors.py index 1303970e7..1857d7d63 100644 --- a/monkey/infection_monkey/payload/ransomware/file_selectors.py +++ b/monkey/infection_monkey/payload/ransomware/file_selectors.py @@ -2,10 +2,10 @@ import filecmp from pathlib import Path from typing import Iterable, Set +from common.utils.file_utils import get_all_regular_files_in_directory from infection_monkey.utils.dir_utils import ( file_extension_filter, filter_files, - get_all_regular_files_in_directory, is_not_shortcut_filter, is_not_symlink_filter, ) diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index da0a5e2e4..9cd20a316 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -2,10 +2,6 @@ from pathlib import Path from typing import Callable, Iterable, Set -def get_all_regular_files_in_directory(dir_path: Path) -> Iterable[Path]: - return filter_files(dir_path.iterdir(), [lambda f: f.is_file()]) - - def filter_files( files: Iterable[Path], file_filters: Iterable[Callable[[Path], bool]] ) -> Iterable[Path]: diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index 48fb95bfc..60c3c1370 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -2,6 +2,7 @@ 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.server_utils.file_utils import create_secure_directory from . import IFileStorageService @@ -50,5 +51,5 @@ class DirectoryFileStorageService(IFileStorageService): return self._storage_directory / safe_file_name def delete_all_files(self): - for file in filter(lambda f: f.is_file(), self._storage_directory.iterdir()): + for file in get_all_regular_files_in_directory(self._storage_directory): file.unlink() diff --git a/monkey/tests/unit_tests/common/utils/test_common_file_utils.py b/monkey/tests/unit_tests/common/utils/test_common_file_utils.py index 79d00d027..aac13839e 100644 --- a/monkey/tests/unit_tests/common/utils/test_common_file_utils.py +++ b/monkey/tests/unit_tests/common/utils/test_common_file_utils.py @@ -1,8 +1,14 @@ import os import pytest +from tests.utils import add_files_to_dir, add_subdirs_to_dir -from common.utils.file_utils import InvalidPath, expand_path, get_file_sha256_hash +from common.utils.file_utils import ( + InvalidPath, + expand_path, + get_all_regular_files_in_directory, + get_file_sha256_hash, +) def test_expand_user(patched_home_env): @@ -26,3 +32,32 @@ def test_expand_path__empty_path_provided(): def test_get_file_sha256_hash(stable_file, stable_file_sha256_hash): assert get_file_sha256_hash(stable_file) == stable_file_sha256_hash + + +SUBDIRS = ["subdir1", "subdir2"] +FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg"] + + +def test_get_all_regular_files_in_directory__no_files(tmp_path, monkeypatch): + add_subdirs_to_dir(tmp_path, SUBDIRS) + + expected_return_value = [] + assert list(get_all_regular_files_in_directory(tmp_path)) == expected_return_value + + +def test_get_all_regular_files_in_directory__has_files(tmp_path, monkeypatch): + add_subdirs_to_dir(tmp_path, SUBDIRS) + files = add_files_to_dir(tmp_path, FILES) + + expected_return_value = sorted(files) + assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value + + +def test_get_all_regular_files_in_directory__subdir_has_files(tmp_path, monkeypatch): + subdirs = add_subdirs_to_dir(tmp_path, SUBDIRS) + add_files_to_dir(subdirs[0], FILES) + + files = add_files_to_dir(tmp_path, FILES) + + expected_return_value = sorted(files) + assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py index adf18bf5a..fe7b499bf 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,66 +1,22 @@ import os import pytest -from tests.utils import is_user_admin +from tests.utils import add_files_to_dir, is_user_admin +from common.utils.file_utils import get_all_regular_files_in_directory from infection_monkey.utils.dir_utils import ( file_extension_filter, filter_files, - get_all_regular_files_in_directory, is_not_shortcut_filter, is_not_symlink_filter, ) SHORTCUT = "shortcut.lnk" FILES = ["file.jpg.zip", "file.xyz", "1.tar", "2.tgz", "2.png", "2.mpg", SHORTCUT] -SUBDIRS = ["subdir1", "subdir2"] - - -def add_subdirs_to_dir(parent_dir): - subdirs = [parent_dir / s for s in SUBDIRS] - - for subdir in subdirs: - subdir.mkdir() - - return subdirs - - -def add_files_to_dir(parent_dir): - files = [parent_dir / f for f in FILES] - - for f in files: - f.touch() - - return files - - -def test_get_all_regular_files_in_directory__no_files(tmp_path, monkeypatch): - add_subdirs_to_dir(tmp_path) - - expected_return_value = [] - assert list(get_all_regular_files_in_directory(tmp_path)) == expected_return_value - - -def test_get_all_regular_files_in_directory__has_files(tmp_path, monkeypatch): - add_subdirs_to_dir(tmp_path) - files = add_files_to_dir(tmp_path) - - expected_return_value = sorted(files) - assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value - - -def test_get_all_regular_files_in_directory__subdir_has_files(tmp_path, monkeypatch): - subdirs = add_subdirs_to_dir(tmp_path) - add_files_to_dir(subdirs[0]) - - files = add_files_to_dir(tmp_path) - - expected_return_value = sorted(files) - assert sorted(get_all_regular_files_in_directory(tmp_path)) == expected_return_value def test_filter_files__no_results(tmp_path): - add_files_to_dir(tmp_path) + add_files_to_dir(tmp_path, FILES) files_in_dir = get_all_regular_files_in_directory(tmp_path) filtered_files = list(filter_files(files_in_dir, [lambda _: False])) @@ -69,7 +25,7 @@ def test_filter_files__no_results(tmp_path): def test_filter_files__all_true(tmp_path): - files = add_files_to_dir(tmp_path) + files = add_files_to_dir(tmp_path, FILES) expected_return_value = sorted(files) files_in_dir = get_all_regular_files_in_directory(tmp_path) @@ -79,7 +35,7 @@ def test_filter_files__all_true(tmp_path): def test_filter_files__multiple_filters(tmp_path): - files = add_files_to_dir(tmp_path) + files = add_files_to_dir(tmp_path, FILES) expected_return_value = sorted(files[4:6]) files_in_dir = get_all_regular_files_in_directory(tmp_path) @@ -93,7 +49,7 @@ def test_filter_files__multiple_filters(tmp_path): def test_file_extension_filter(tmp_path): valid_extensions = {".zip", ".xyz"} - files = add_files_to_dir(tmp_path) + files = add_files_to_dir(tmp_path, FILES) files_in_dir = get_all_regular_files_in_directory(tmp_path) filtered_files = filter_files(files_in_dir, [file_extension_filter(valid_extensions)]) @@ -105,7 +61,7 @@ def test_file_extension_filter(tmp_path): os.name == "nt" and not is_user_admin(), reason="Test requires admin rights on Windows" ) def test_is_not_symlink_filter(tmp_path): - files = add_files_to_dir(tmp_path) + files = add_files_to_dir(tmp_path, FILES) link_path = tmp_path / "symlink.test" link_path.symlink_to(files[0], target_is_directory=False) @@ -118,7 +74,7 @@ def test_is_not_symlink_filter(tmp_path): def test_is_not_shortcut_filter(tmp_path): - add_files_to_dir(tmp_path) + add_files_to_dir(tmp_path, FILES) files_in_dir = get_all_regular_files_in_directory(tmp_path) filtered_files = list(filter_files(files_in_dir, [is_not_shortcut_filter])) diff --git a/monkey/tests/utils.py b/monkey/tests/utils.py index 9b57a9cc7..b2268e572 100644 --- a/monkey/tests/utils.py +++ b/monkey/tests/utils.py @@ -1,5 +1,7 @@ import ctypes import os +from pathlib import Path +from typing import Iterable def is_user_admin(): @@ -11,3 +13,21 @@ def is_user_admin(): def raise_(ex): raise ex + + +def add_subdirs_to_dir(parent_dir: Path, subdirs: Iterable[str]) -> Iterable[Path]: + subdir_paths = [parent_dir / s for s in subdirs] + + for subdir in subdir_paths: + subdir.mkdir() + + return subdir_paths + + +def add_files_to_dir(parent_dir: Path, file_names: Iterable[str]) -> Iterable[Path]: + files = [parent_dir / f for f in file_names] + + for f in files: + f.touch() + + return files From 379a71d8e2da7f364aee7471ac14c35f9cbcb716 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 14:44:31 -0400 Subject: [PATCH 12/35] Common: Add DIContainer --- monkey/common/__init__.py | 1 + monkey/common/di_container.py | 83 ++++++++ .../unit_tests/common/test_di_container.py | 178 ++++++++++++++++++ 3 files changed, 262 insertions(+) create mode 100644 monkey/common/di_container.py create mode 100644 monkey/tests/unit_tests/common/test_di_container.py diff --git a/monkey/common/__init__.py b/monkey/common/__init__.py index e69de29bb..5f2cba2f1 100644 --- a/monkey/common/__init__.py +++ b/monkey/common/__init__.py @@ -0,0 +1 @@ +from .di_container import DIContainer diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py new file mode 100644 index 000000000..1f5a2e03c --- /dev/null +++ b/monkey/common/di_container.py @@ -0,0 +1,83 @@ +import inspect +from typing import Any, MutableMapping, Type, TypeVar + +T = TypeVar("T") + + +class DIContainer: + """ + A dependency injection (DI) container that uses type annotations to resolve and inject + dependencies. + """ + + def __init__(self): + self._type_registry = {} + self._instance_registry = {} + + def register(self, interface: Type[T], concrete_type: Type[T]): + """ + Register a concrete type that satisfies a given interface. + + :param interface: An interface or abstract base class that other classes depend upon + :param concrete_type: A type (class) that implements `interface` + """ + self._type_registry[interface] = concrete_type + DIContainer._del_key(self._instance_registry, interface) + + def register_instance(self, interface: Type[T], instance: T): + """ + Register a concrete instance that satisfies a given interface. + + :param interface: An interface or abstract base class that other classes depend upon + :param instance: An instance (object) of a type that implements `interface` + """ + self._instance_registry[interface] = instance + DIContainer._del_key(self._type_registry, interface) + + def resolve(self, type_: Type[T]) -> T: + """ + Resolves all dependencies and returns a new instance of `type_` using constructor dependency + injection. + + :param type_: A type (class) to construct. + :return: An instance of `type_` + """ + args = [] + + # TODO: Need to handle keyword-only arguments, defaults, varars, etc. + for arg_type in inspect.getfullargspec(type_).annotations.values(): + new_instance = self._resolve_instance(arg_type) + args.append(new_instance) + + return type_(*args) + + def _resolve_instance(self, arg_type: Type): + if arg_type in self._type_registry: + return self._type_registry[arg_type]() + elif arg_type in self._instance_registry: + return self._instance_registry[arg_type] + + raise ValueError(f'Failed to resolve unknown type "{arg_type.__name__}"') + + def release(self, interface: Type[T]): + """ + Deregister's an interface + + :param interface: The interface to release + """ + DIContainer._del_key(self._type_registry, interface) + DIContainer._del_key(self._instance_registry, interface) + + @staticmethod + def _del_key(mapping: MutableMapping[T, Any], key: T): + """ + Deletes key from mapping. Unlike the `del` keyword, this function does not raise a KeyError + if the key does not exist. + + :param MutableMapping: A mapping from which a key will be deleted + :param key: A key to delete from `mapping` + """ + try: + del mapping[key] + except KeyError: + pass diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py new file mode 100644 index 000000000..d3b6abf57 --- /dev/null +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -0,0 +1,178 @@ +import pytest + +from common import DIContainer + + +class IServiceA: + pass + + +class IServiceB: + pass + + +class ServiceA(IServiceA): + pass + + +class ServiceB(IServiceB): + pass + + +class TestClass1: + __test__ = False + + def __init__(self, service_a: IServiceA): + self.service_a = service_a + + +class TestClass2: + __test__ = False + + def __init__(self, service_b: IServiceB): + self.service_b = service_b + + +class TestClass3: + __test__ = False + + def __init__(self, service_a: IServiceA, service_b: IServiceB): + self.service_a = service_a + self.service_b = service_b + + +@pytest.fixture +def container(): + return DIContainer() + + +def test_register_resolve(container): + container.register(IServiceA, ServiceA) + test_1 = container.resolve(TestClass1) + + assert isinstance(test_1.service_a, ServiceA) + + +def test_correct_instance_type_injected(container): + container.register(IServiceA, ServiceA) + container.register(IServiceB, ServiceB) + test_1 = container.resolve(TestClass1) + test_2 = container.resolve(TestClass2) + + assert isinstance(test_1.service_a, ServiceA) + assert isinstance(test_2.service_b, ServiceB) + + +def test_multiple_correct_instance_types_injected(container): + container.register(IServiceA, ServiceA) + container.register(IServiceB, ServiceB) + test_3 = container.resolve(TestClass3) + + assert isinstance(test_3.service_a, ServiceA) + assert isinstance(test_3.service_b, ServiceB) + + +def test_register_instance(container): + service_a_instance = ServiceA() + + container.register_instance(IServiceA, service_a_instance) + test_1 = container.resolve(TestClass1) + + assert id(service_a_instance) == id(test_1.service_a) + + +def test_register_multiple_instances(container): + service_a_instance = ServiceA() + service_b_instance = ServiceB() + + container.register_instance(IServiceA, service_a_instance) + container.register_instance(IServiceB, service_b_instance) + test_3 = container.resolve(TestClass3) + + assert id(service_a_instance) == id(test_3.service_a) + assert id(service_b_instance) == id(test_3.service_b) + + +def test_register_mixed_instance_and_type(container): + service_a_instance = ServiceA() + + container.register_instance(IServiceA, service_a_instance) + container.register(IServiceB, ServiceB) + test_2 = container.resolve(TestClass2) + test_3 = container.resolve(TestClass3) + + assert id(service_a_instance) == id(test_3.service_a) + assert isinstance(test_2.service_b, ServiceB) + assert isinstance(test_3.service_b, ServiceB) + assert id(test_2.service_b) != id(test_3.service_b) + + +def test_unregistered_type(): + container = DIContainer() + with pytest.raises(ValueError): + container.resolve(TestClass1) + + +def test_type_registration_overwritten(container): + class ServiceA2(IServiceA): + pass + + container.register(IServiceA, ServiceA) + container.register(IServiceA, ServiceA2) + test_1 = container.resolve(TestClass1) + + assert isinstance(test_1.service_a, ServiceA2) + + +def test_instance_registration_overwritten(container): + service_a_instance_1 = ServiceA() + service_a_instance_2 = ServiceA() + + container.register_instance(IServiceA, service_a_instance_1) + container.register_instance(IServiceA, service_a_instance_2) + test_1 = container.resolve(TestClass1) + + assert id(test_1.service_a) != id(service_a_instance_1) + assert id(test_1.service_a) == id(service_a_instance_2) + + +def test_type_overrides_instance(container): + service_a_instance = ServiceA() + + container.register_instance(IServiceA, service_a_instance) + container.register(IServiceA, ServiceA) + test_1 = container.resolve(TestClass1) + + assert id(test_1.service_a) != id(service_a_instance) + assert isinstance(test_1.service_a, ServiceA) + + +def test_instance_overrides_type(container): + service_a_instance = ServiceA() + + container.register(IServiceA, ServiceA) + container.register_instance(IServiceA, service_a_instance) + test_1 = container.resolve(TestClass1) + + assert id(test_1.service_a) == id(service_a_instance) + + +def test_release_type(): + container = DIContainer() + + container.register(IServiceA, ServiceA) + container.release(IServiceA) + + with pytest.raises(ValueError): + container.resolve(TestClass1) + + +def test_release_instance(): + container = DIContainer() + service_a_instance = ServiceA() + + container.register_instance(IServiceA, service_a_instance) + container.release(IServiceA) + + with pytest.raises(ValueError): + container.resolve(TestClass1) From 435b619a5d8b8f6edb513f8c41a991a7a53041c4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 25 Apr 2022 15:56:59 -0400 Subject: [PATCH 13/35] Common: Add recursive resolution to DIContainer --- monkey/common/di_container.py | 23 +++++-- .../unit_tests/common/test_di_container.py | 69 ++++++++++++++++--- 2 files changed, 78 insertions(+), 14 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 1f5a2e03c..c36de1eba 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -44,21 +44,32 @@ class DIContainer: """ args = [] - # TODO: Need to handle keyword-only arguments, defaults, varars, etc. + # TODO: Need to handle keyword-only arguments, defaults, varargs, etc. for arg_type in inspect.getfullargspec(type_).annotations.values(): - new_instance = self._resolve_instance(arg_type) - args.append(new_instance) + instance = self._resolve_arg_type(arg_type) + args.append(instance) return type_(*args) - def _resolve_instance(self, arg_type: Type): + def _resolve_arg_type(self, arg_type: Type[T]) -> T: if arg_type in self._type_registry: - return self._type_registry[arg_type]() + return self._resolve_type(arg_type) elif arg_type in self._instance_registry: - return self._instance_registry[arg_type] + return self._resolve_instance(arg_type) raise ValueError(f'Failed to resolve unknown type "{arg_type.__name__}"') + def _resolve_type(self, arg_type: Type[T]) -> T: + try: + return self._type_registry[arg_type]() + except TypeError: + # arg_type has dependencies that must be resolved. Recursively call resolve() to + # construct an instance of arg_type with all of the requesite dependencies injected. + return self.resolve(self._type_registry[arg_type]) + + def _resolve_instance(self, arg_type: Type[T]) -> T: + return self._instance_registry[arg_type] + def release(self, interface: Type[T]): """ Deregister's an interface diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index d3b6abf57..63cfcfd13 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -1,13 +1,15 @@ +import abc + import pytest from common import DIContainer -class IServiceA: +class IServiceA(metaclass=abc.ABCMeta): pass -class IServiceB: +class IServiceB(metaclass=abc.ABCMeta): pass @@ -157,9 +159,7 @@ def test_instance_overrides_type(container): assert id(test_1.service_a) == id(service_a_instance) -def test_release_type(): - container = DIContainer() - +def test_release_type(container): container.register(IServiceA, ServiceA) container.release(IServiceA) @@ -167,12 +167,65 @@ def test_release_type(): container.resolve(TestClass1) -def test_release_instance(): - container = DIContainer() +def test_release_instance(container): service_a_instance = ServiceA() - container.register_instance(IServiceA, service_a_instance) + container.release(IServiceA) with pytest.raises(ValueError): container.resolve(TestClass1) + + +class IServiceC(metaclass=abc.ABCMeta): + pass + + +class ServiceC(IServiceC): + def __init__(self, service_a: IServiceA): + self.service_a = service_a + + +class TestClass4: + def __init__(self, service_c: IServiceC): + self.service_c = service_c + + +def test_recursive_resolution__depth_2(container): + service_a_instance = ServiceA() + container.register_instance(IServiceA, service_a_instance) + container.register(IServiceC, ServiceC) + + test4 = container.resolve(TestClass4) + + assert isinstance(test4.service_c, ServiceC) + assert id(test4.service_c.service_a) == id(service_a_instance) + + +class IServiceD(metaclass=abc.ABCMeta): + pass + + +class ServiceD(IServiceD): + def __init__(self, service_c: IServiceC, service_b: IServiceB): + self.service_b = service_b + self.service_c = service_c + + +class TestClass5: + def __init__(self, service_d: IServiceD): + self.service_d = service_d + + +def test_recursive_resolution__depth_3(container): + container.register(IServiceA, ServiceA) + container.register(IServiceB, ServiceB) + container.register(IServiceC, ServiceC) + container.register(IServiceD, ServiceD) + + test5 = container.resolve(TestClass5) + + assert isinstance(test5.service_d, ServiceD) + assert isinstance(test5.service_d.service_b, ServiceB) + assert isinstance(test5.service_d.service_c, ServiceC) + assert isinstance(test5.service_d.service_c.service_a, ServiceA) From 7a62434364e7dd9cda92aba81c64b003e281e8a2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 01:21:36 -0400 Subject: [PATCH 14/35] Common: Resolve registered instances and types directly --- monkey/common/di_container.py | 23 +++++++++------- .../unit_tests/common/test_di_container.py | 27 ++++++++++++++++--- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index c36de1eba..8ee7fcb1f 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -42,24 +42,29 @@ class DIContainer: :param type_: A type (class) to construct. :return: An instance of `type_` """ + try: + return self._resolve_type(type_) + except ValueError: + pass + args = [] # TODO: Need to handle keyword-only arguments, defaults, varargs, etc. for arg_type in inspect.getfullargspec(type_).annotations.values(): - instance = self._resolve_arg_type(arg_type) + instance = self._resolve_type(arg_type) args.append(instance) return type_(*args) - def _resolve_arg_type(self, arg_type: Type[T]) -> T: - if arg_type in self._type_registry: - return self._resolve_type(arg_type) - elif arg_type in self._instance_registry: - return self._resolve_instance(arg_type) + def _resolve_type(self, type_: Type[T]) -> T: + if type_ in self._type_registry: + return self._construct_new_instance(type_) + elif type_ in self._instance_registry: + return self._retrieve_registered_instance(type_) - raise ValueError(f'Failed to resolve unknown type "{arg_type.__name__}"') + raise ValueError(f'Failed to resolve unknown type "{type_.__name__}"') - def _resolve_type(self, arg_type: Type[T]) -> T: + def _construct_new_instance(self, arg_type: Type[T]) -> T: try: return self._type_registry[arg_type]() except TypeError: @@ -67,7 +72,7 @@ class DIContainer: # construct an instance of arg_type with all of the requesite dependencies injected. return self.resolve(self._type_registry[arg_type]) - def _resolve_instance(self, arg_type: Type[T]) -> T: + def _retrieve_registered_instance(self, arg_type: Type[T]) -> T: return self._instance_registry[arg_type] def release(self, interface: Type[T]): diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index 63cfcfd13..dab1a8f29 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -6,7 +6,9 @@ from common import DIContainer class IServiceA(metaclass=abc.ABCMeta): - pass + @abc.abstractmethod + def do_something(self): + pass class IServiceB(metaclass=abc.ABCMeta): @@ -14,7 +16,8 @@ class IServiceB(metaclass=abc.ABCMeta): class ServiceA(IServiceA): - pass + def do_something(self): + pass class ServiceB(IServiceB): @@ -117,7 +120,8 @@ def test_unregistered_type(): def test_type_registration_overwritten(container): class ServiceA2(IServiceA): - pass + def do_something(self): + pass container.register(IServiceA, ServiceA) container.register(IServiceA, ServiceA2) @@ -229,3 +233,20 @@ def test_recursive_resolution__depth_3(container): assert isinstance(test5.service_d.service_b, ServiceB) assert isinstance(test5.service_d.service_c, ServiceC) assert isinstance(test5.service_d.service_c.service_a, ServiceA) + + +def test_resolve_registered_interface(container): + container.register(IServiceA, ServiceA) + + resolved_instance = container.resolve(IServiceA) + + assert isinstance(resolved_instance, ServiceA) + + +def test_resolve_registered_instance(container): + service_a_instance = ServiceA() + container.register_instance(IServiceA, service_a_instance) + + service_a_actual_instance = container.resolve(IServiceA) + + assert id(service_a_actual_instance) == id(service_a_instance) From e78bffb414c8438c2d43a0106253d5812ce2fdab Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 01:27:17 -0400 Subject: [PATCH 15/35] Common: Add note about varargs and kwargs to resolve() docstring --- monkey/common/di_container.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 8ee7fcb1f..4ad246ec7 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -37,7 +37,8 @@ class DIContainer: def resolve(self, type_: Type[T]) -> T: """ Resolves all dependencies and returns a new instance of `type_` using constructor dependency - injection. + injection. Note that only positional arguments are resolved. Varargs, keyword-only args, and + default values are ignored. :param type_: A type (class) to construct. :return: An instance of `type_` @@ -49,7 +50,6 @@ class DIContainer: args = [] - # TODO: Need to handle keyword-only arguments, defaults, varargs, etc. for arg_type in inspect.getfullargspec(type_).annotations.values(): instance = self._resolve_type(arg_type) args.append(instance) From 7382407be07c53dbd80dbc5a94e79d42778ef468 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 01:37:41 -0400 Subject: [PATCH 16/35] Common: Add DIContainer.resolve_dependencies() --- monkey/common/di_container.py | 18 +++++++++++++++--- .../unit_tests/common/test_di_container.py | 10 ++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 4ad246ec7..9642ffb1f 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, MutableMapping, Type, TypeVar +from typing import Any, MutableMapping, Sequence, Type, TypeVar T = TypeVar("T") @@ -40,7 +40,7 @@ class DIContainer: injection. Note that only positional arguments are resolved. Varargs, keyword-only args, and default values are ignored. - :param type_: A type (class) to construct. + :param type_: A type (class) to construct :return: An instance of `type_` """ try: @@ -48,13 +48,25 @@ class DIContainer: except ValueError: pass + args = self.resolve_dependencies(type_) + return type_(*args) + + def resolve_dependencies(self, type_: Type[T]) -> Sequence[Any]: + """ + Resolves all dependencies of type_ and returns a Sequence of objects that correspond type_'s + dependencies. Note that only positional arguments are resolved. Varargs, keyword-only args, + and default values are ignored. + + :param type_: A type (class) to resolve dependencies for + :return: An Sequence of dependencies to be injected into type_'s constructor + """ args = [] for arg_type in inspect.getfullargspec(type_).annotations.values(): instance = self._resolve_type(arg_type) args.append(instance) - return type_(*args) + return tuple(args) def _resolve_type(self, type_: Type[T]) -> T: if type_ in self._type_registry: diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index dab1a8f29..1d1e4ea5c 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -250,3 +250,13 @@ def test_resolve_registered_instance(container): service_a_actual_instance = container.resolve(IServiceA) assert id(service_a_actual_instance) == id(service_a_instance) + + +def test_resolve_dependencies(container): + container.register(IServiceA, ServiceA) + container.register(IServiceB, ServiceB) + + dependencies = container.resolve_dependencies(TestClass3) + + assert isinstance(dependencies[0], ServiceA) + assert isinstance(dependencies[1], ServiceB) From c16705241a45d7566042b355d5f79c44ff7c76c6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 01:41:20 -0400 Subject: [PATCH 17/35] Common: Use a more precise, custom UnregisteredTypeError in DIContainer --- monkey/common/di_container.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 9642ffb1f..55ba31923 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -4,6 +4,10 @@ from typing import Any, MutableMapping, Sequence, Type, TypeVar T = TypeVar("T") +class UnregisteredTypeError(ValueError): + pass + + class DIContainer: """ A dependency injection (DI) container that uses type annotations to resolve and inject @@ -45,7 +49,7 @@ class DIContainer: """ try: return self._resolve_type(type_) - except ValueError: + except UnregisteredTypeError: pass args = self.resolve_dependencies(type_) @@ -74,7 +78,7 @@ class DIContainer: elif type_ in self._instance_registry: return self._retrieve_registered_instance(type_) - raise ValueError(f'Failed to resolve unknown type "{type_.__name__}"') + raise UnregisteredTypeError(f'Failed to resolve unregistered type "{type_.__name__}"') def _construct_new_instance(self, arg_type: Type[T]) -> T: try: From 3d931b11f99a571d1823367bfc1b0421907dc49d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 02:18:43 -0400 Subject: [PATCH 18/35] Island: Fix incorrect type hint in AuthenticationService --- .../cc/services/authentication/authentication_service.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/authentication/authentication_service.py b/monkey/monkey_island/cc/services/authentication/authentication_service.py index ddb096215..93fbf32f4 100644 --- a/monkey/monkey_island/cc/services/authentication/authentication_service.py +++ b/monkey/monkey_island/cc/services/authentication/authentication_service.py @@ -1,3 +1,5 @@ +from pathlib import Path + import bcrypt from common.utils.exceptions import ( @@ -23,7 +25,7 @@ class AuthenticationService: # static/singleton hybrids. At the moment, this requires invasive refactoring that's # not a priority. @classmethod - def initialize(cls, data_dir: str, user_datastore: IUserDatastore): + def initialize(cls, data_dir: Path, user_datastore: IUserDatastore): cls.DATA_DIR = data_dir cls.user_datastore = user_datastore From 9b52f3f21b7cfc30e7b47091197a153c2a8dcc7e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 02:19:55 -0400 Subject: [PATCH 19/35] Island: Add missing type hint to LocalMonkeyRunService constructor --- monkey/monkey_island/cc/services/run_local_monkey.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/run_local_monkey.py b/monkey/monkey_island/cc/services/run_local_monkey.py index be08352e8..a1ae3f6e3 100644 --- a/monkey/monkey_island/cc/services/run_local_monkey.py +++ b/monkey/monkey_island/cc/services/run_local_monkey.py @@ -3,6 +3,7 @@ import os import platform import stat import subprocess +from pathlib import Path from shutil import copyfile from monkey_island.cc.resources.monkey_download import get_agent_executable_path @@ -19,7 +20,7 @@ class LocalMonkeyRunService: # static/singleton hybrids. At the moment, this requires invasive refactoring that's # not a priority. @classmethod - def initialize(cls, data_dir): + def initialize(cls, data_dir: Path): cls.DATA_DIR = data_dir @staticmethod From 8c3477a000bb95827d59d621a22f8b00f332be1a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 02:32:02 -0400 Subject: [PATCH 20/35] Common: Raise TypeError if DIContainer.register() called with instance --- monkey/common/di_container.py | 7 +++++++ monkey/tests/unit_tests/common/test_di_container.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 55ba31923..6086440e7 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -25,6 +25,13 @@ class DIContainer: :param interface: An interface or abstract base class that other classes depend upon :param concrete_type: A type (class) that implements `interface` """ + if not inspect.isclass(concrete_type): + raise TypeError( + "Expected a class, but received an instance of type " + f'"{concrete_type.__class__.__name__}"; Pass a class, not an instance, to ' + "register(), or use register_instance() instead." + ) + self._type_registry[interface] = concrete_type DIContainer._del_key(self._instance_registry, interface) diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index 1d1e4ea5c..1caa55ca4 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -260,3 +260,9 @@ def test_resolve_dependencies(container): assert isinstance(dependencies[0], ServiceA) assert isinstance(dependencies[1], ServiceB) + + +def test_register_instance_as_type(container): + service_a_instance = ServiceA() + with pytest.raises(TypeError): + container.register(IServiceA, service_a_instance) From d084b003678b8f39742dabd15f66d92fa14db2ec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 02:36:17 -0400 Subject: [PATCH 21/35] Island: Use DIContainer to construct custom PBA resources --- monkey/monkey_island/cc/app.py | 109 ++++++++++-------- monkey/monkey_island/cc/server_setup.py | 15 ++- .../monkey_island/cc/services/initialize.py | 17 ++- 3 files changed, 86 insertions(+), 55 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index beb9e1468..da8d312f7 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,12 +1,13 @@ import os import uuid from datetime import timedelta -from pathlib import Path +from typing import Type import flask_restful from flask import Flask, Response, send_from_directory from werkzeug.exceptions import NotFound +from common import DIContainer from monkey_island.cc.database import database, mongo from monkey_island.cc.resources.agent_controls import StopAgentCheck, StopAllAgents from monkey_island.cc.resources.attack.attack_report import AttackReport @@ -47,7 +48,6 @@ from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFinding from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.custom_json_encoder import CustomJSONEncoder -from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService from monkey_island.cc.services.representations import output_json @@ -113,73 +113,90 @@ def init_app_url_rules(app): app.add_url_rule("/", "serve_static_file", serve_static_file) -def init_api_resources(api, data_dir: Path): - api.add_resource(Root, "/api") - api.add_resource(Registration, "/api/registration") - api.add_resource(Authenticate, "/api/auth") - api.add_resource( +# TODO: Come up with a better name than FlaskResourceManager +class FlaskResourceManager: + def __init__(self, api: flask_restful.Api, container: DIContainer): + self._api = api + self._container = container + + def add_resource(self, resource: Type[flask_restful.Resource], *urls: str): + dependencies = self._container.resolve_dependencies(resource) + self._api.add_resource(resource, *urls, resource_class_args=dependencies) + + +def init_api_resources(flask_resource_manager: FlaskResourceManager): + flask_resource_manager.add_resource(Root, "/api") + flask_resource_manager.add_resource(Registration, "/api/registration") + flask_resource_manager.add_resource(Authenticate, "/api/auth") + flask_resource_manager.add_resource( Monkey, "/api/agent", "/api/agent/", "/api/agent//", ) - api.add_resource(LocalRun, "/api/local-monkey") - api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/") + flask_resource_manager.add_resource(LocalRun, "/api/local-monkey") + flask_resource_manager.add_resource( + Telemetry, "/api/telemetry", "/api/telemetry/" + ) - api.add_resource(IslandMode, "/api/island-mode") - api.add_resource(IslandConfiguration, "/api/configuration/island") - api.add_resource(ConfigurationExport, "/api/configuration/export") - api.add_resource(ConfigurationImport, "/api/configuration/import") - api.add_resource( + flask_resource_manager.add_resource(IslandMode, "/api/island-mode") + flask_resource_manager.add_resource(IslandConfiguration, "/api/configuration/island") + flask_resource_manager.add_resource(ConfigurationExport, "/api/configuration/export") + flask_resource_manager.add_resource(ConfigurationImport, "/api/configuration/import") + flask_resource_manager.add_resource( MonkeyDownload, "/api/agent/download/", ) - api.add_resource(NetMap, "/api/netmap") - api.add_resource(Edge, "/api/netmap/edge") - api.add_resource(Node, "/api/netmap/node") - api.add_resource(NodeStates, "/api/netmap/node-states") + flask_resource_manager.add_resource(NetMap, "/api/netmap") + flask_resource_manager.add_resource(Edge, "/api/netmap/edge") + flask_resource_manager.add_resource(Node, "/api/netmap/node") + flask_resource_manager.add_resource(NodeStates, "/api/netmap/node-states") - api.add_resource(SecurityReport, "/api/report/security") - api.add_resource(ZeroTrustReport, "/api/report/zero-trust/") - api.add_resource(AttackReport, "/api/report/attack") - api.add_resource(RansomwareReport, "/api/report/ransomware") - api.add_resource(ManualExploitation, "/api/exploitations/manual") - api.add_resource(MonkeyExploitation, "/api/exploitations/monkey") + flask_resource_manager.add_resource(SecurityReport, "/api/report/security") + flask_resource_manager.add_resource( + ZeroTrustReport, "/api/report/zero-trust/" + ) + flask_resource_manager.add_resource(AttackReport, "/api/report/attack") + flask_resource_manager.add_resource(RansomwareReport, "/api/report/ransomware") + flask_resource_manager.add_resource(ManualExploitation, "/api/exploitations/manual") + flask_resource_manager.add_resource(MonkeyExploitation, "/api/exploitations/monkey") - api.add_resource(ZeroTrustFindingEvent, "/api/zero-trust/finding-event/") - api.add_resource(TelemetryFeed, "/api/telemetry-feed") - api.add_resource(Log, "/api/log") - api.add_resource(IslandLog, "/api/log/island/download") + flask_resource_manager.add_resource( + ZeroTrustFindingEvent, "/api/zero-trust/finding-event/" + ) + flask_resource_manager.add_resource(TelemetryFeed, "/api/telemetry-feed") + flask_resource_manager.add_resource(Log, "/api/log") + flask_resource_manager.add_resource(IslandLog, "/api/log/island/download") - # This is temporary until we get DI all worked out. - file_storage_service = DirectoryFileStorageService(data_dir / "custom_pbas") - api.add_resource( + flask_resource_manager.add_resource( PBAFileDownload, "/api/pba/download/", - resource_class_kwargs={"file_storage_service": file_storage_service}, ) - api.add_resource( + flask_resource_manager.add_resource( FileUpload, "/api/file-upload/", "/api/file-upload/?load=", "/api/file-upload/?restore=", - resource_class_kwargs={"file_storage_service": file_storage_service}, ) - api.add_resource(PropagationCredentials, "/api/propagation-credentials/") - api.add_resource(RemoteRun, "/api/remote-monkey") - api.add_resource(VersionUpdate, "/api/version-update") - api.add_resource(StopAgentCheck, "/api/monkey-control/needs-to-stop/") - api.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") + flask_resource_manager.add_resource( + PropagationCredentials, "/api/propagation-credentials/" + ) + flask_resource_manager.add_resource(RemoteRun, "/api/remote-monkey") + flask_resource_manager.add_resource(VersionUpdate, "/api/version-update") + flask_resource_manager.add_resource( + StopAgentCheck, "/api/monkey-control/needs-to-stop/" + ) + flask_resource_manager.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") # Resources used by black box tests - api.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") - api.add_resource(ClearCaches, "/api/test/clear-caches") - api.add_resource(LogBlackboxEndpoint, "/api/test/log") - api.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") + flask_resource_manager.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") + flask_resource_manager.add_resource(ClearCaches, "/api/test/clear-caches") + flask_resource_manager.add_resource(LogBlackboxEndpoint, "/api/test/log") + flask_resource_manager.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") -def init_app(mongo_url, data_dir: Path): +def init_app(mongo_url: str, container: DIContainer): app = Flask(__name__) api = flask_restful.Api(app) @@ -188,6 +205,8 @@ def init_app(mongo_url, data_dir: Path): init_app_config(app, mongo_url) init_app_services(app) init_app_url_rules(app) - init_api_resources(api, data_dir) + + flask_resource_manager = FlaskResourceManager(api, container) + init_api_resources(flask_resource_manager) return app diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 072cb32a8..df41bb137 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -16,6 +16,7 @@ MONKEY_ISLAND_DIR_BASE_PATH = str(Path(__file__).parent.parent) if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH) +from common import DIContainer # noqa: E402 from common.version import get_version # noqa: E402 from monkey_island.cc.app import init_app # noqa: E402 from monkey_island.cc.arg_parser import IslandCmdArgs # noqa: E402 @@ -47,7 +48,7 @@ def run_monkey_island(): _exit_on_invalid_config_options(config_options) _configure_logging(config_options) - _initialize_globals(config_options.data_dir) + container = _initialize_globals(config_options.data_dir) mongo_db_process = None if config_options.start_mongodb: @@ -56,7 +57,7 @@ def run_monkey_island(): _connect_to_mongodb(mongo_db_process) _configure_gevent_exception_handling(config_options.data_dir) - _start_island_server(island_args.setup_only, config_options) + _start_island_server(island_args.setup_only, config_options, container) def _extract_config(island_args: IslandCmdArgs) -> IslandConfigOptions: @@ -88,8 +89,8 @@ def _configure_logging(config_options): setup_logging(config_options.data_dir, config_options.log_level) -def _initialize_globals(data_dir: Path): - initialize_services(data_dir) +def _initialize_globals(data_dir: Path) -> DIContainer: + return initialize_services(data_dir) def _start_mongodb(data_dir: Path) -> MongoDbProcess: @@ -127,9 +128,11 @@ def _configure_gevent_exception_handling(data_dir): hub.handle_error = GeventHubErrorHandler(hub, logger) -def _start_island_server(should_setup_only, config_options: IslandConfigOptions): +def _start_island_server( + should_setup_only: bool, config_options: IslandConfigOptions, container: DIContainer +): populate_exporter_list() - app = init_app(mongo_setup.MONGO_URL, Path(config_options.data_dir)) + app = init_app(mongo_setup.MONGO_URL, container) if should_setup_only: logger.warning("Setup only flag passed. Exiting.") diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index f3ff10395..4a4b2e4af 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,13 +1,22 @@ -from monkey_island.cc.services import DirectoryFileStorageService +from pathlib import Path + +from common import DIContainer +from monkey_island.cc.services import DirectoryFileStorageService, IFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService from . import AuthenticationService, JsonFileUserDatastore -def initialize_services(data_dir): - # This is temporary until we get DI all worked out. - PostBreachFilesService.initialize(DirectoryFileStorageService(data_dir / "custom_pbas")) +def initialize_services(data_dir: Path) -> DIContainer: + container = DIContainer() + container.register_instance( + IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas") + ) + # This is temporary until we get DI all worked out. + PostBreachFilesService.initialize(container.resolve(IFileStorageService)) LocalMonkeyRunService.initialize(data_dir) AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) + + return container From e296cd5225b676ebdb2315c4796085c132ee1de8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 12:58:22 -0400 Subject: [PATCH 22/35] UT: Refactor PBA resource upload/download tests to use DI --- .../monkey_island/cc/resources/conftest.py | 24 +++++++-- .../cc/resources/test_pba_file_download.py | 53 ++++++++++++++++--- .../cc/resources/test_pba_file_upload.py | 45 ++++++++++++++-- 3 files changed, 107 insertions(+), 15 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 903528be2..db58ed173 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,3 +1,5 @@ +from unittest.mock import MagicMock + import flask_jwt_extended import flask_restful import pytest @@ -10,14 +12,27 @@ from monkey_island.cc.services.representations import output_json @pytest.fixture -def flask_client(monkeypatch_session, tmp_path): +def flask_client(monkeypatch_session): monkeypatch_session.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) - with mock_init_app(tmp_path).test_client() as client: + container = MagicMock() + container.resolve_dependencies.return_value = [] + + with mock_init_app(container).test_client() as client: yield client -def mock_init_app(data_dir): +@pytest.fixture +def build_flask_client(monkeypatch_session): + def inner(container): + monkeypatch_session.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) + + return mock_init_app(container).test_client() + + return inner + + +def mock_init_app(container): app = Flask(__name__) app.config["SECRET_KEY"] = "test_key" @@ -25,7 +40,8 @@ def mock_init_app(data_dir): api.representations = {"application/json": output_json} monkey_island.cc.app.init_app_url_rules(app) - monkey_island.cc.app.init_api_resources(api, data_dir) + flask_resource_manager = monkey_island.cc.app.FlaskResourceManager(api, container) + monkey_island.cc.app.init_api_resources(flask_resource_manager) flask_jwt_extended.JWTManager(app) 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 5b6dcad3f..6b76115cb 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,17 +1,54 @@ -def test_file_download_endpoint(tmp_path, flask_client): - file_contents = "HelloWorld!" - file_name = "test_file" - (tmp_path / "custom_pbas" / file_name).write_text(file_contents) +import io +from typing import BinaryIO - resp = flask_client.get(f"/api/pba/download/{file_name}") +import pytest + +from common import DIContainer +from monkey_island.cc.services import IFileStorageService + +FILE_NAME = "test_file" +FILE_CONTENTS = b"HelloWorld!" + + +class MockFileStorageService(IFileStorageService): + 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 OSError() + + 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, tmp_path): + container = DIContainer() + container.register(IFileStorageService, MockFileStorageService) + + with build_flask_client(container) as flask_client: + yield flask_client + + +def test_file_download_endpoint(tmp_path, flask_client): + resp = flask_client.get(f"/api/pba/download/{FILE_NAME}") assert resp.status_code == 200 - assert next(resp.response).decode() == file_contents + assert next(resp.response) == FILE_CONTENTS def test_file_download_endpoint_404(tmp_path, flask_client): - file_name = "nonexistant_file" + nonexistant_file_name = "nonexistant_file" - resp = flask_client.get(f"/api/pba/download/{file_name}") + resp = flask_client.get(f"/api/pba/download/{nonexistant_file_name}") assert resp.status_code == 404 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 ca5f0f175..28b908c86 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 @@ -1,9 +1,16 @@ +import io +from typing import BinaryIO + import pytest from tests.utils import raise_ +from common import DIContainer from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE +from monkey_island.cc.services import IFileStorageService -TEST_FILE = b"""-----------------------------1 +TEST_FILE_CONTENTS = b"m0nk3y" +TEST_FILE = ( + b"""-----------------------------1 Content-Disposition: form-data; name="filepond" {} @@ -11,8 +18,11 @@ Content-Disposition: form-data; name="filepond" Content-Disposition: form-data; name="filepond"; filename="test.py" Content-Type: text/x-python -m0nk3y +""" + + TEST_FILE_CONTENTS + + b""" -----------------------------1--""" +) @pytest.fixture @@ -29,6 +39,35 @@ def mock_get_config_value(monkeypatch): ) +class MockFileStorageService(IFileStorageService): + def __init__(self): + self._file = None + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._file = io.BytesIO(file_contents.read()) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + if self._file is None: + # TODO: Add FileRetrievalError + raise OSError() + return self._file + + def delete_file(self, unsafe_file_name: str): + self._file = None + + def delete_all_files(self): + self.delete_file("") + + +@pytest.fixture +def flask_client(build_flask_client, tmp_path): + container = DIContainer() + container.register(IFileStorageService, MockFileStorageService) + + with build_flask_client(container) as flask_client: + yield flask_client + + @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, mock_set_config_value): resp = flask_client.post( @@ -89,7 +128,7 @@ def test_pba_file_upload_endpoint( resp_get = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") assert resp_get.status_code == 200 - assert resp_get.data.decode() == "m0nk3y" + assert resp_get.data == TEST_FILE_CONTENTS # Closing the response closes the file handle, else it can't be deleted resp_get.close() From 1d938e58f844c008dfc5611a13f421e97675f0cd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 13:38:52 -0400 Subject: [PATCH 23/35] Island: Prefix private functions with "_" in FileUpload resource --- .../monkey_island/cc/resources/pba_file_upload.py | 13 ++++++------- .../cc/resources/test_pba_file_upload.py | 2 +- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index bb2608872..5fdde2045 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -30,7 +30,6 @@ class FileUpload(flask_restful.Resource): # TODO: Fix references/coupling to filepond # TODO: Replace "file_type" with "target_os" or similar - # TODO: Prefix private functions with "_" # TODO: Add comment explaining why this is basically a duplicate of the endpoint in the # PBAFileDownload resource. @jwt_required @@ -40,7 +39,7 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file to send, linux or windows :return: Returns file contents """ - if self.is_pba_file_type_supported(file_type): + if self._is_pba_file_type_supported(file_type): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") # Verify that file_name is indeed a file from config @@ -67,10 +66,10 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file was received, linux or windows :return: Returns flask response object with uploaded file's filename """ - if self.is_pba_file_type_supported(file_type): + if self._is_pba_file_type_supported(file_type): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") - filename = self.upload_pba_file( + filename = self._upload_pba_file( # TODO: This "filepond" string can be changed to be more generic in the `react-filepond` # component. request.files["filepond"], @@ -80,7 +79,7 @@ class FileUpload(flask_restful.Resource): response = Response(response=filename, status=200, mimetype="text/plain") return response - def upload_pba_file(self, file_storage: FileStorage, is_linux=True): + def _upload_pba_file(self, file_storage: FileStorage, is_linux=True): """ Uploads PBA file to island's file system :param request_: Request object containing PBA file @@ -103,7 +102,7 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file was deleted, linux of windows :return: Empty response """ - if self.is_pba_file_type_supported(file_type): + if self._is_pba_file_type_supported(file_type): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") filename_path = ( @@ -117,5 +116,5 @@ class FileUpload(flask_restful.Resource): return make_response({}, 200) @staticmethod - def is_pba_file_type_supported(file_type: str) -> bool: + def _is_pba_file_type_supported(file_type: str) -> bool: return file_type not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE} 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 28b908c86..00e10757d 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 @@ -94,7 +94,7 @@ def test_pba_file_upload_post__internal_server_error( flask_client, pba_os, monkeypatch, mock_set_config_value ): monkeypatch.setattr( - "monkey_island.cc.resources.pba_file_upload.FileUpload.upload_pba_file", + "monkey_island.cc.resources.pba_file_upload.FileUpload._upload_pba_file", lambda x, y: raise_(Exception()), ) From 55c4f68902222482d98061e2452b85ad01bdfb96 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 13:41:16 -0400 Subject: [PATCH 24/35] Island: Rename file_type -> target_os in FileUpload resource --- monkey/monkey_island/cc/app.py | 6 ++-- .../cc/resources/pba_file_upload.py | 29 +++++++++---------- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index da8d312f7..404f5731b 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -174,9 +174,9 @@ def init_api_resources(flask_resource_manager: FlaskResourceManager): ) flask_resource_manager.add_resource( FileUpload, - "/api/file-upload/", - "/api/file-upload/?load=", - "/api/file-upload/?restore=", + "/api/file-upload/", + "/api/file-upload/?load=", + "/api/file-upload/?restore=", ) flask_resource_manager.add_resource( diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 5fdde2045..d7b170a9f 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -29,21 +29,20 @@ class FileUpload(flask_restful.Resource): self._file_storage_service = file_storage_service # TODO: Fix references/coupling to filepond - # TODO: Replace "file_type" with "target_os" or similar # TODO: Add comment explaining why this is basically a duplicate of the endpoint in the # PBAFileDownload resource. @jwt_required - def get(self, file_type): + def get(self, target_os): """ Sends file to filepond - :param file_type: Type indicates which file to send, linux or windows + :param target_os: Indicates which file to send, linux or windows :return: Returns file contents """ - if self._is_pba_file_type_supported(file_type): + if self._is_target_os_supported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") # Verify that file_name is indeed a file from config - if file_type == LINUX_PBA_TYPE: + if target_os == LINUX_PBA_TYPE: # TODO: Make these paths Tuples so we don't need to copy them filename = ConfigService.get_config_value(copy.deepcopy(PBA_LINUX_FILENAME_PATH)) else: @@ -60,20 +59,20 @@ class FileUpload(flask_restful.Resource): return make_response({"error": error_msg}, 404) @jwt_required - def post(self, file_type): + def post(self, target_os): """ Receives user's uploaded file from filepond - :param file_type: Type indicates which file was received, linux or windows + :param target_os: Type indicates which file was received, linux or windows :return: Returns flask response object with uploaded file's filename """ - if self._is_pba_file_type_supported(file_type): + if self._is_target_os_supported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") filename = self._upload_pba_file( # TODO: This "filepond" string can be changed to be more generic in the `react-filepond` # component. request.files["filepond"], - (file_type == LINUX_PBA_TYPE), + (target_os == LINUX_PBA_TYPE), ) response = Response(response=filename, status=200, mimetype="text/plain") @@ -96,17 +95,17 @@ class FileUpload(flask_restful.Resource): return filename @jwt_required - def delete(self, file_type): + def delete(self, target_os): """ Deletes file that has been deleted on the front end - :param file_type: Type indicates which file was deleted, linux of windows + :param target_os: Type indicates which file was deleted, linux of windows :return: Empty response """ - if self._is_pba_file_type_supported(file_type): + if self._is_target_os_supported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") filename_path = ( - PBA_LINUX_FILENAME_PATH if file_type == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH + PBA_LINUX_FILENAME_PATH if target_os == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH ) filename = ConfigService.get_config_value(filename_path) if filename: @@ -116,5 +115,5 @@ class FileUpload(flask_restful.Resource): return make_response({}, 200) @staticmethod - def _is_pba_file_type_supported(file_type: str) -> bool: - return file_type not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE} + def _is_target_os_supported(target_os: str) -> bool: + return target_os not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE} From fefdd9f024031caf79b039a19755217bb077f10f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 13:49:17 -0400 Subject: [PATCH 25/35] UT: Remove mocking of private method in test_pba_file_upload.py --- .../cc/resources/test_pba_file_upload.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) 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 00e10757d..fa9e02b5c 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 @@ -60,9 +60,14 @@ class MockFileStorageService(IFileStorageService): @pytest.fixture -def flask_client(build_flask_client, tmp_path): +def file_storage_service(): + return MockFileStorageService() + + +@pytest.fixture +def flask_client(build_flask_client, file_storage_service): container = DIContainer() - container.register(IFileStorageService, MockFileStorageService) + container.register_instance(IFileStorageService, file_storage_service) with build_flask_client(container) as flask_client: yield flask_client @@ -91,12 +96,9 @@ def test_pba_file_upload_post__invalid(flask_client, monkeypatch, mock_set_confi @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_post__internal_server_error( - flask_client, pba_os, monkeypatch, mock_set_config_value + flask_client, pba_os, mock_set_config_value, file_storage_service ): - monkeypatch.setattr( - "monkey_island.cc.resources.pba_file_upload.FileUpload._upload_pba_file", - lambda x, y: raise_(Exception()), - ) + file_storage_service.save_file = lambda x, y: raise_(Exception()) resp = flask_client.post( f"/api/file-upload/{pba_os}", From 98033c96c04930ee66b4ff3c66ac50553439fceb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 13:50:22 -0400 Subject: [PATCH 26/35] UT: Remove unnecessary monkeypatch fixtures from test_pba_file_upload.py --- .../cc/resources/test_pba_file_upload.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) 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 fa9e02b5c..6f615e19f 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 @@ -74,7 +74,7 @@ def flask_client(build_flask_client, file_storage_service): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, mock_set_config_value): +def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): resp = flask_client.post( f"/api/file-upload/{pba_os}", data=TEST_FILE, @@ -84,7 +84,7 @@ def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, mock_set_config assert resp.status_code == 200 -def test_pba_file_upload_post__invalid(flask_client, monkeypatch, mock_set_config_value): +def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): resp = flask_client.post( "/api/file-upload/bogus", data=TEST_FILE, @@ -110,16 +110,14 @@ def test_pba_file_upload_post__internal_server_error( @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_get__file_not_found( - flask_client, pba_os, monkeypatch, mock_get_config_value -): +def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_config_value): resp = flask_client.get(f"/api/file-upload/{pba_os}?load=bogus_mogus.py") assert resp.status_code == 404 @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_endpoint( - flask_client, pba_os, monkeypatch, mock_get_config_value, mock_set_config_value + flask_client, pba_os, mock_get_config_value, mock_set_config_value ): resp_post = flask_client.post( f"/api/file-upload/{pba_os}", @@ -146,7 +144,7 @@ def test_pba_file_upload_endpoint( def test_pba_file_upload_endpoint__invalid( - flask_client, monkeypatch, mock_set_config_value, mock_get_config_value + flask_client, mock_set_config_value, mock_get_config_value ): resp_post = flask_client.post( "/api/file-upload/bogus", From 54d34d18160630dc9f1917620e12651d9c6820b2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:19:24 -0400 Subject: [PATCH 27/35] Island: Rename FlaskResourceManager -> FlaskDIWrapper --- monkey/monkey_island/cc/app.py | 89 ++++++++----------- .../monkey_island/cc/resources/conftest.py | 2 +- 2 files changed, 40 insertions(+), 51 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 404f5731b..7cb34a5a3 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -113,8 +113,7 @@ def init_app_url_rules(app): app.add_url_rule("/", "serve_static_file", serve_static_file) -# TODO: Come up with a better name than FlaskResourceManager -class FlaskResourceManager: +class FlaskDIWrapper: def __init__(self, api: flask_restful.Api, container: DIContainer): self._api = api self._container = container @@ -124,76 +123,66 @@ class FlaskResourceManager: self._api.add_resource(resource, *urls, resource_class_args=dependencies) -def init_api_resources(flask_resource_manager: FlaskResourceManager): - flask_resource_manager.add_resource(Root, "/api") - flask_resource_manager.add_resource(Registration, "/api/registration") - flask_resource_manager.add_resource(Authenticate, "/api/auth") - flask_resource_manager.add_resource( +def init_api_resources(api: FlaskDIWrapper): + api.add_resource(Root, "/api") + api.add_resource(Registration, "/api/registration") + api.add_resource(Authenticate, "/api/auth") + api.add_resource( Monkey, "/api/agent", "/api/agent/", "/api/agent//", ) - flask_resource_manager.add_resource(LocalRun, "/api/local-monkey") - flask_resource_manager.add_resource( - Telemetry, "/api/telemetry", "/api/telemetry/" - ) + api.add_resource(LocalRun, "/api/local-monkey") + api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/") - flask_resource_manager.add_resource(IslandMode, "/api/island-mode") - flask_resource_manager.add_resource(IslandConfiguration, "/api/configuration/island") - flask_resource_manager.add_resource(ConfigurationExport, "/api/configuration/export") - flask_resource_manager.add_resource(ConfigurationImport, "/api/configuration/import") - flask_resource_manager.add_resource( + api.add_resource(IslandMode, "/api/island-mode") + api.add_resource(IslandConfiguration, "/api/configuration/island") + api.add_resource(ConfigurationExport, "/api/configuration/export") + api.add_resource(ConfigurationImport, "/api/configuration/import") + api.add_resource( MonkeyDownload, "/api/agent/download/", ) - flask_resource_manager.add_resource(NetMap, "/api/netmap") - flask_resource_manager.add_resource(Edge, "/api/netmap/edge") - flask_resource_manager.add_resource(Node, "/api/netmap/node") - flask_resource_manager.add_resource(NodeStates, "/api/netmap/node-states") + api.add_resource(NetMap, "/api/netmap") + api.add_resource(Edge, "/api/netmap/edge") + api.add_resource(Node, "/api/netmap/node") + api.add_resource(NodeStates, "/api/netmap/node-states") - flask_resource_manager.add_resource(SecurityReport, "/api/report/security") - flask_resource_manager.add_resource( - ZeroTrustReport, "/api/report/zero-trust/" - ) - flask_resource_manager.add_resource(AttackReport, "/api/report/attack") - flask_resource_manager.add_resource(RansomwareReport, "/api/report/ransomware") - flask_resource_manager.add_resource(ManualExploitation, "/api/exploitations/manual") - flask_resource_manager.add_resource(MonkeyExploitation, "/api/exploitations/monkey") + api.add_resource(SecurityReport, "/api/report/security") + api.add_resource(ZeroTrustReport, "/api/report/zero-trust/") + api.add_resource(AttackReport, "/api/report/attack") + api.add_resource(RansomwareReport, "/api/report/ransomware") + api.add_resource(ManualExploitation, "/api/exploitations/manual") + api.add_resource(MonkeyExploitation, "/api/exploitations/monkey") - flask_resource_manager.add_resource( - ZeroTrustFindingEvent, "/api/zero-trust/finding-event/" - ) - flask_resource_manager.add_resource(TelemetryFeed, "/api/telemetry-feed") - flask_resource_manager.add_resource(Log, "/api/log") - flask_resource_manager.add_resource(IslandLog, "/api/log/island/download") + api.add_resource(ZeroTrustFindingEvent, "/api/zero-trust/finding-event/") + api.add_resource(TelemetryFeed, "/api/telemetry-feed") + api.add_resource(Log, "/api/log") + api.add_resource(IslandLog, "/api/log/island/download") - flask_resource_manager.add_resource( + api.add_resource( PBAFileDownload, "/api/pba/download/", ) - flask_resource_manager.add_resource( + api.add_resource( FileUpload, "/api/file-upload/", "/api/file-upload/?load=", "/api/file-upload/?restore=", ) - flask_resource_manager.add_resource( - PropagationCredentials, "/api/propagation-credentials/" - ) - flask_resource_manager.add_resource(RemoteRun, "/api/remote-monkey") - flask_resource_manager.add_resource(VersionUpdate, "/api/version-update") - flask_resource_manager.add_resource( - StopAgentCheck, "/api/monkey-control/needs-to-stop/" - ) - flask_resource_manager.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") + api.add_resource(PropagationCredentials, "/api/propagation-credentials/") + api.add_resource(RemoteRun, "/api/remote-monkey") + api.add_resource(VersionUpdate, "/api/version-update") + api.add_resource(StopAgentCheck, "/api/monkey-control/needs-to-stop/") + api.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") # Resources used by black box tests - flask_resource_manager.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") - flask_resource_manager.add_resource(ClearCaches, "/api/test/clear-caches") - flask_resource_manager.add_resource(LogBlackboxEndpoint, "/api/test/log") - flask_resource_manager.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") + api.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") + api.add_resource(ClearCaches, "/api/test/clear-caches") + api.add_resource(LogBlackboxEndpoint, "/api/test/log") + api.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") def init_app(mongo_url: str, container: DIContainer): @@ -206,7 +195,7 @@ def init_app(mongo_url: str, container: DIContainer): init_app_services(app) init_app_url_rules(app) - flask_resource_manager = FlaskResourceManager(api, container) + flask_resource_manager = FlaskDIWrapper(api, container) init_api_resources(flask_resource_manager) return app 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 db58ed173..723c777b8 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -40,7 +40,7 @@ def mock_init_app(container): api.representations = {"application/json": output_json} monkey_island.cc.app.init_app_url_rules(app) - flask_resource_manager = monkey_island.cc.app.FlaskResourceManager(api, container) + flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) monkey_island.cc.app.init_api_resources(flask_resource_manager) flask_jwt_extended.JWTManager(app) From 92349d8f8ed0b3f5209116208d924214cbe5401e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:21:02 -0400 Subject: [PATCH 28/35] UT: Remove collection warnings for TestClass{4,5} --- monkey/tests/unit_tests/common/test_di_container.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index 1caa55ca4..b0109b11d 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -191,6 +191,8 @@ class ServiceC(IServiceC): class TestClass4: + __test__ = False + def __init__(self, service_c: IServiceC): self.service_c = service_c @@ -217,6 +219,8 @@ class ServiceD(IServiceD): class TestClass5: + __test__ = False + def __init__(self, service_d: IServiceD): self.service_d = service_d From 8f7215034d9cba9651b4c7c8b62f1e5d5ddfc6b8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:26:27 -0400 Subject: [PATCH 29/35] Island: Add paranoid check to avoid directory traversal attacks --- .../cc/services/directory_file_storage_service.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index 60c3c1370..35e6989f2 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -45,10 +45,13 @@ class DirectoryFileStorageService(IFileStorageService): def _get_safe_file_path(self, unsafe_file_name: str): # Remove any path information from the file name. safe_file_name = Path(unsafe_file_name).resolve().name + safe_file_path = (self._storage_directory / safe_file_name).resolve() - # TODO: Add super paranoid check + # This is a paranoid check to avoid directory traversal attacks. + if self._storage_directory.resolve() not in safe_file_path.parents: + raise ValueError(f"The file named {unsafe_file_name} can not be safely retrieved") - return self._storage_directory / safe_file_name + return safe_file_path def delete_all_files(self): for file in get_all_regular_files_in_directory(self._storage_directory): From a0b4dc1bcbb106ba33279ff4c830d1e94d30898e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:36:50 -0400 Subject: [PATCH 30/35] Island: Ignore exception in delete_file() if file not found --- .../cc/services/directory_file_storage_service.py | 5 ++++- .../cc/services/test_directory_file_storage_service.py | 7 +++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index 35e6989f2..d6fd267fe 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -40,7 +40,10 @@ class DirectoryFileStorageService(IFileStorageService): def delete_file(self, unsafe_file_name: str): safe_file_path = self._get_safe_file_path(unsafe_file_name) - safe_file_path.unlink() + try: + safe_file_path.unlink() + except FileNotFoundError: + pass def _get_safe_file_path(self, unsafe_file_name: str): # Remove any path information from the file name. diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py index 09ed308c5..07f65151d 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py @@ -119,3 +119,10 @@ def test_remove_all_files__skip_directories(tmp_path): for file in tmp_path.iterdir(): assert file.name == test_dir.name + + +def test_remove_nonexistant_file(tmp_path): + fss = DirectoryFileStorageService(tmp_path) + + # This test will fail if this call raises an exception. + fss.delete_file("nonexistant_file.txt") From 4ddcd5e9a820fcc2ef8ce1a4e64cd8dc9f86f285 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:48:22 -0400 Subject: [PATCH 31/35] Island: Raise FileRetrievalError in DirectoryFileStorageService --- .../monkey_island/cc/resources/pba_file_download.py | 6 +++--- monkey/monkey_island/cc/resources/pba_file_upload.py | 6 +++--- monkey/monkey_island/cc/services/__init__.py | 2 +- .../cc/services/directory_file_storage_service.py | 8 ++++++-- .../cc/services/i_file_storage_service.py | 5 +++++ .../cc/resources/test_pba_file_download.py | 4 ++-- .../cc/resources/test_pba_file_upload.py | 5 ++--- .../services/test_directory_file_storage_service.py | 11 +++++++++-- 8 files changed, 31 insertions(+), 16 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index ba5857c53..a11e964b1 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -3,7 +3,7 @@ import logging import flask_restful from flask import make_response, send_file -from monkey_island.cc.services import IFileStorageService +from monkey_island.cc.services import FileRetrievalError, IFileStorageService logger = logging.getLogger(__file__) @@ -23,7 +23,7 @@ class PBAFileDownload(flask_restful.Resource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except OSError as ex: - error_msg = f"Failed to open file {filename}: {ex}" + except FileRetrievalError as err: + error_msg = f"Failed to open file {filename}: {err}" logger.error(error_msg) return make_response({"error": error_msg}, 404) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index d7b170a9f..2a27b324a 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -9,7 +9,7 @@ from werkzeug.utils import secure_filename from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.services import IFileStorageService +from monkey_island.cc.services import FileRetrievalError, IFileStorageService from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__file__) @@ -53,8 +53,8 @@ class FileUpload(flask_restful.Resource): # `send_file()` handles the closing of the open file. return send_file(file, mimetype="application/octet-stream") - except OSError as ex: - error_msg = f"Failed to open file {filename}: {ex}" + except FileRetrievalError as err: + error_msg = f"Failed to open file {filename}: {err}" logger.error(error_msg) return make_response({"error": error_msg}, 404) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 3f57b4fc0..43aa39382 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,4 +1,4 @@ -from .i_file_storage_service import IFileStorageService +from .i_file_storage_service import IFileStorageService, FileRetrievalError from .directory_file_storage_service import DirectoryFileStorageService from .authentication.authentication_service import AuthenticationService diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index d6fd267fe..d2d0c811d 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -5,7 +5,7 @@ from typing import BinaryIO from common.utils.file_utils import get_all_regular_files_in_directory from monkey_island.cc.server_utils.file_utils import create_secure_directory -from . import IFileStorageService +from . import FileRetrievalError, IFileStorageService class DirectoryFileStorageService(IFileStorageService): @@ -35,7 +35,11 @@ class DirectoryFileStorageService(IFileStorageService): def open_file(self, unsafe_file_name: str) -> BinaryIO: safe_file_path = self._get_safe_file_path(unsafe_file_name) - return open(safe_file_path, "rb") + + try: + return open(safe_file_path, "rb") + except OSError as err: + 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/monkey_island/cc/services/i_file_storage_service.py b/monkey/monkey_island/cc/services/i_file_storage_service.py index a27cbf6ad..5903c07d1 100644 --- a/monkey/monkey_island/cc/services/i_file_storage_service.py +++ b/monkey/monkey_island/cc/services/i_file_storage_service.py @@ -2,6 +2,10 @@ import abc from typing import BinaryIO +class FileRetrievalError(ValueError): + pass + + class IFileStorageService(metaclass=abc.ABCMeta): """ A service that allows the storage and retrieval of individual files. @@ -25,6 +29,7 @@ class IFileStorageService(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 :rtype: io.BinaryIO + :raises FileRetrievalError: if the file cannot be opened """ 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 6b76115cb..d0a8ec48f 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 @@ -4,7 +4,7 @@ from typing import BinaryIO import pytest from common import DIContainer -from monkey_island.cc.services import IFileStorageService +from monkey_island.cc.services import FileRetrievalError, IFileStorageService FILE_NAME = "test_file" FILE_CONTENTS = b"HelloWorld!" @@ -19,7 +19,7 @@ class MockFileStorageService(IFileStorageService): def open_file(self, unsafe_file_name: str) -> BinaryIO: if unsafe_file_name != FILE_NAME: - raise OSError() + raise FileRetrievalError() return self._file 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 6f615e19f..da1eb2c02 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 @@ -6,7 +6,7 @@ from tests.utils import raise_ from common import DIContainer from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE -from monkey_island.cc.services import IFileStorageService +from monkey_island.cc.services import FileRetrievalError, IFileStorageService TEST_FILE_CONTENTS = b"m0nk3y" TEST_FILE = ( @@ -48,8 +48,7 @@ class MockFileStorageService(IFileStorageService): def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: - # TODO: Add FileRetrievalError - raise OSError() + raise FileRetrievalError() return self._file def delete_file(self, unsafe_file_name: str): diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py index 07f65151d..ea40aa49d 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_directory_file_storage_service.py @@ -5,10 +5,10 @@ import pytest from tests.monkey_island.utils import assert_linux_permissions, assert_windows_permissions from monkey_island.cc.server_utils.file_utils import is_windows_os -from monkey_island.cc.services import DirectoryFileStorageService +from monkey_island.cc.services import DirectoryFileStorageService, FileRetrievalError -def test_error_if_file(tmp_path): +def test_error_if_storage_directory_is_file(tmp_path): new_file = tmp_path / "new_file.txt" new_file.write_text("HelloWorld!") @@ -126,3 +126,10 @@ def test_remove_nonexistant_file(tmp_path): # This test will fail if this call raises an exception. fss.delete_file("nonexistant_file.txt") + + +def test_open_nonexistant_file(tmp_path): + fss = DirectoryFileStorageService(tmp_path) + + with pytest.raises(FileRetrievalError): + fss.open_file("nonexistant_file.txt") From e8e879091d1b215ffabb358fa25b8e70a270afd1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 19:54:12 -0400 Subject: [PATCH 32/35] Island: Add logging to DirectoryFileStorageService --- .../cc/services/directory_file_storage_service.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index d2d0c811d..f8aeebe21 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -1,3 +1,4 @@ +import logging import shutil from pathlib import Path from typing import BinaryIO @@ -7,6 +8,8 @@ from monkey_island.cc.server_utils.file_utils import create_secure_directory from . import FileRetrievalError, IFileStorageService +logger = logging.getLogger(__name__) + class DirectoryFileStorageService(IFileStorageService): """ @@ -30,6 +33,7 @@ class DirectoryFileStorageService(IFileStorageService): def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): safe_file_path = self._get_safe_file_path(unsafe_file_name) + logger.debug(f"Saving file contents to {safe_file_path}") with open(safe_file_path, "wb") as dest: shutil.copyfileobj(file_contents, dest) @@ -37,14 +41,17 @@ class DirectoryFileStorageService(IFileStorageService): safe_file_path = self._get_safe_file_path(unsafe_file_name) try: + logger.debug(f"Opening {safe_file_path}") return open(safe_file_path, "rb") except OSError as err: + logger.error(err) 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) try: + logger.debug(f"Deleting {safe_file_path}") safe_file_path.unlink() except FileNotFoundError: pass @@ -58,8 +65,10 @@ class DirectoryFileStorageService(IFileStorageService): if self._storage_directory.resolve() not in safe_file_path.parents: raise ValueError(f"The file named {unsafe_file_name} can not be safely retrieved") + logger.debug(f"Unsafe file name {unsafe_file_name} sanitized: {safe_file_path}") return safe_file_path def delete_all_files(self): for file in get_all_regular_files_in_directory(self._storage_directory): + logger.debug(f"Deleting {file}") file.unlink() From 77753aca3f6f7a5a4f3234eb9d78da67eeca1044 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 27 Apr 2022 09:07:27 -0400 Subject: [PATCH 33/35] Island: Rename _initialize_globals() -> _initialize_di_container() --- monkey/monkey_island/cc/server_setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index df41bb137..399b8ebc9 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -48,7 +48,7 @@ def run_monkey_island(): _exit_on_invalid_config_options(config_options) _configure_logging(config_options) - container = _initialize_globals(config_options.data_dir) + container = _initialize_di_container(config_options.data_dir) mongo_db_process = None if config_options.start_mongodb: @@ -89,7 +89,7 @@ def _configure_logging(config_options): setup_logging(config_options.data_dir, config_options.log_level) -def _initialize_globals(data_dir: Path) -> DIContainer: +def _initialize_di_container(data_dir: Path) -> DIContainer: return initialize_services(data_dir) From 1e9627470c894728ce6adac0d99c51e4824f0540 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 27 Apr 2022 09:18:55 -0400 Subject: [PATCH 34/35] Common: Check that injected types are subtypes of interfaces --- monkey/common/di_container.py | 13 ++++++++++++- monkey/tests/unit_tests/common/test_di_container.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/monkey/common/di_container.py b/monkey/common/di_container.py index 6086440e7..8559e8036 100644 --- a/monkey/common/di_container.py +++ b/monkey/common/di_container.py @@ -29,7 +29,12 @@ class DIContainer: raise TypeError( "Expected a class, but received an instance of type " f'"{concrete_type.__class__.__name__}"; Pass a class, not an instance, to ' - "register(), or use register_instance() instead." + "register(), or use register_instance() instead" + ) + + if not issubclass(concrete_type, interface): + raise TypeError( + f'Class "{concrete_type.__name__}" is not a subclass of {interface.__name__}' ) self._type_registry[interface] = concrete_type @@ -42,6 +47,12 @@ class DIContainer: :param interface: An interface or abstract base class that other classes depend upon :param instance: An instance (object) of a type that implements `interface` """ + if not isinstance(instance, interface): + raise TypeError( + f'The provided instance of type "{instance.__class__.__name__}" ' + f"is not an instance of {interface.__name__}" + ) + self._instance_registry[interface] = instance DIContainer._del_key(self._type_registry, interface) diff --git a/monkey/tests/unit_tests/common/test_di_container.py b/monkey/tests/unit_tests/common/test_di_container.py index b0109b11d..f7ea5534c 100644 --- a/monkey/tests/unit_tests/common/test_di_container.py +++ b/monkey/tests/unit_tests/common/test_di_container.py @@ -270,3 +270,14 @@ def test_register_instance_as_type(container): service_a_instance = ServiceA() with pytest.raises(TypeError): container.register(IServiceA, service_a_instance) + + +def test_register_conflicting_type(container): + with pytest.raises(TypeError): + container.register(IServiceA, ServiceB) + + +def test_register_instance_with_conflicting_type(container): + service_b_instance = ServiceB() + with pytest.raises(TypeError): + container.register_instance(IServiceA, service_b_instance) From 3bf0695a667f047451fae737bd88762d67cae3b5 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 27 Apr 2022 09:27:54 -0400 Subject: [PATCH 35/35] Island: Add comment/documentation about delete_file()'s idempotence --- .../cc/services/directory_file_storage_service.py | 1 + monkey/monkey_island/cc/services/i_file_storage_service.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/monkey/monkey_island/cc/services/directory_file_storage_service.py b/monkey/monkey_island/cc/services/directory_file_storage_service.py index f8aeebe21..a4be91522 100644 --- a/monkey/monkey_island/cc/services/directory_file_storage_service.py +++ b/monkey/monkey_island/cc/services/directory_file_storage_service.py @@ -54,6 +54,7 @@ class DirectoryFileStorageService(IFileStorageService): logger.debug(f"Deleting {safe_file_path}") safe_file_path.unlink() except FileNotFoundError: + # This method is idempotent. pass def _get_safe_file_path(self, unsafe_file_name: str): diff --git a/monkey/monkey_island/cc/services/i_file_storage_service.py b/monkey/monkey_island/cc/services/i_file_storage_service.py index 5903c07d1..31c49cd00 100644 --- a/monkey/monkey_island/cc/services/i_file_storage_service.py +++ b/monkey/monkey_island/cc/services/i_file_storage_service.py @@ -38,6 +38,9 @@ class IFileStorageService(metaclass=abc.ABCMeta): """ Delete a file + This method will delete the file specified by `unsafe_file_name`. This operation is + idempotent and will succeed if the file to be deleted does not exist. + :param unsafe_file_name: An unsanitized file name that identifies the file to be deleted """ pass