From b4ced896b64985a7d6713986c3db87b74d60356a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 10:21:33 -0400 Subject: [PATCH 1/6] UT: Improve formatting in test_pba_file_upload_endpoint() --- .../monkey_island/cc/resources/test_pba_file_upload.py | 3 +-- 1 file changed, 1 insertion(+), 2 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 e4cb5a487..d49d293bd 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 @@ -129,10 +129,9 @@ def test_pba_file_upload_endpoint( resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;") resp_get_del = flask_client.get(url_with_filename) + assert resp_post.status_code == 200 - assert resp_delete.status_code == 200 - assert resp_get_del.status_code == 404 From 2e7bcd54df15b920474cbac36735ea2a481675ec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 10:22:21 -0400 Subject: [PATCH 2/6] Island: Remove ConfigService from PBA FileUpload resource --- .../cc/resources/pba_file_upload.py | 55 +++++++++++++------ .../cc/resources/test_pba_file_upload.py | 44 +++++---------- 2 files changed, 54 insertions(+), 45 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 4deef93c3..fab3deb36 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -1,15 +1,14 @@ import logging +from dataclasses import replace from http import HTTPStatus from flask import Response, make_response, request, send_file from werkzeug.utils import secure_filename as sanitize_filename -from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH from monkey_island.cc import repository -from monkey_island.cc.repository import IFileRepository +from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required -from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__file__) @@ -30,8 +29,13 @@ class FileUpload(AbstractResource): "/api/file-upload/?restore=", ] - def __init__(self, file_storage_repository: IFileRepository): + def __init__( + self, + file_storage_repository: IFileRepository, + agent_configuration_repository: IAgentConfigurationRepository, + ): self._file_storage_service = file_storage_repository + self._agent_configuration_repository = agent_configuration_repository # This endpoint is basically a duplicate of PBAFileDownload.get(). They serve slightly different # purposes. This endpoint is authenticated, whereas the one in PBAFileDownload can not be (at @@ -48,11 +52,12 @@ class FileUpload(AbstractResource): if self._is_target_os_supported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") + agent_configuration = self._agent_configuration_repository.get_configuration() # Verify that file_name is indeed a file from config if target_os == LINUX_PBA_TYPE: - filename = ConfigService.get_config_value(PBA_LINUX_FILENAME_PATH) + filename = agent_configuration.custom_pbas.linux_filename else: - filename = ConfigService.get_config_value(PBA_WINDOWS_FILENAME_PATH) + filename = agent_configuration.custom_pbas.windows_filename try: file = self._file_storage_service.open_file(filename) @@ -77,15 +82,27 @@ class FileUpload(AbstractResource): safe_filename = sanitize_filename(file_storage.filename) self._file_storage_service.save_file(safe_filename, file_storage.stream) - ConfigService.set_config_value( - (PBA_LINUX_FILENAME_PATH if target_os == LINUX_PBA_TYPE else PBA_WINDOWS_FILENAME_PATH), - safe_filename, - ) + try: + self._update_config(target_os, safe_filename) + except Exception as err: + self._file_storage.delete_file(safe_filename) + raise err # API Spec: HTTP status code should be 201 response = Response(response=safe_filename, status=200, mimetype="text/plain") return response + def _update_config(self, target_os: str, safe_filename: str): + agent_configuration = self._agent_configuration_repository.get_configuration() + + if target_os == LINUX_PBA_TYPE: + custom_pbas = replace(agent_configuration.custom_pbas, linux_filename=safe_filename) + else: + custom_pbas = replace(agent_configuration.custom_pbas, windows_filename=safe_filename) + + updated_agent_configuration = replace(agent_configuration, custom_pbas=custom_pbas) + self._agent_configuration_repository.store_configuration(updated_agent_configuration) + @jwt_required def delete(self, target_os): """ @@ -96,13 +113,19 @@ class FileUpload(AbstractResource): if self._is_target_os_supported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") - filename_path = ( - PBA_LINUX_FILENAME_PATH if target_os == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH - ) - filename = ConfigService.get_config_value(filename_path) - if filename: + original_agent_configuration = self._agent_configuration_repository.get_configuration() + self._update_config(target_os, "") + + if target_os == LINUX_PBA_TYPE: + filename = original_agent_configuration.custom_pbas.linux_filename + else: + filename = original_agent_configuration.custom_pbas.windows_filename + + try: self._file_storage_service.delete_file(filename) - ConfigService.set_config_value(filename_path, "") + except Exception as err: + self._agent_configuration_repository.store_configuration(original_agent_configuration) + raise err # API Spec: HTTP status code should be 204 return make_response({}, 200) 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 d49d293bd..423dd262c 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,10 +1,10 @@ import pytest from tests.common import StubDIContainer -from tests.monkey_island import SingleFileRepository +from tests.monkey_island import InMemoryAgentConfigurationRepository, SingleFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ -from monkey_island.cc.repository import IFileRepository +from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload TEST_FILE_CONTENTS = b"m0nk3y" @@ -24,36 +24,28 @@ Content-Type: text/x-python ) -@pytest.fixture -def mock_set_config_value(monkeypatch): - monkeypatch.setattr( - "monkey_island.cc.services.config.ConfigService.set_config_value", lambda _, __: None - ) - - -@pytest.fixture -def mock_get_config_value(monkeypatch): - monkeypatch.setattr( - "monkey_island.cc.services.config.ConfigService.get_config_value", lambda _: "test.py" - ) - - @pytest.fixture def file_repository(): return SingleFileRepository() @pytest.fixture -def flask_client(build_flask_client, file_repository): +def agent_configuration_repository(): + return InMemoryAgentConfigurationRepository() + + +@pytest.fixture +def flask_client(build_flask_client, file_repository, agent_configuration_repository): container = StubDIContainer() container.register_instance(IFileRepository, file_repository) + container.register_instance(IAgentConfigurationRepository, agent_configuration_repository) 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, mock_set_config_value): +def test_pba_file_upload_post(flask_client, pba_os): url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( url, @@ -64,7 +56,7 @@ def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): assert resp.status_code == 200 -def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): +def test_pba_file_upload_post__invalid(flask_client): url = get_url_for_resource(FileUpload, target_os="bogus") resp = flask_client.post( url, @@ -76,9 +68,7 @@ def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post__internal_server_error( - flask_client, pba_os, mock_set_config_value, file_repository -): +def test_pba_file_upload_post__internal_server_error(flask_client, pba_os, file_repository): file_repository.save_file = lambda x, y: raise_(Exception()) url = get_url_for_resource(FileUpload, target_os=pba_os) @@ -92,7 +82,7 @@ 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, mock_get_config_value): +def test_pba_file_upload_get__file_not_found(flask_client, pba_os): url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") resp = flask_client.get(url) assert resp.status_code == 404 @@ -108,9 +98,7 @@ def test_file_download_endpoint_500(open_error_flask_client, pba_os): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_endpoint( - flask_client, pba_os, mock_get_config_value, mock_set_config_value -): +def test_pba_file_upload_endpoint(flask_client, pba_os): url_with_os = get_url_for_resource(FileUpload, target_os=pba_os) resp_post = flask_client.post( @@ -135,9 +123,7 @@ def test_pba_file_upload_endpoint( assert resp_get_del.status_code == 404 -def test_pba_file_upload_endpoint__invalid( - flask_client, mock_set_config_value, mock_get_config_value -): +def test_pba_file_upload_endpoint__invalid(flask_client): url_with_os = get_url_for_resource(FileUpload, target_os="bogus") resp_post = flask_client.post( From 0d45c5fb3ebee6457a3c0fc63959b7d987d2596d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 10:26:03 -0400 Subject: [PATCH 3/6] Island: Add notes and comments to PBA FileUpload resource --- monkey/monkey_island/cc/resources/pba_file_upload.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index fab3deb36..eee10a3d9 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -17,6 +17,7 @@ LINUX_PBA_TYPE = "PBAlinux" WINDOWS_PBA_TYPE = "PBAwindows" +# NOTE: This resource will be reworked when the Custom PBA feature is rebuilt as a payload plugin. class FileUpload(AbstractResource): # API Spec: FileUpload -> PBAFileUpload. Change endpoint accordingly. """ @@ -37,6 +38,9 @@ class FileUpload(AbstractResource): self._file_storage_service = file_storage_repository self._agent_configuration_repository = agent_configuration_repository + # NOTE: None of these methods are thread-safe. Don't forget to fix that when this becomes a + # payload plugin. + # This endpoint is basically a duplicate of PBAFileDownload.get(). They serve slightly different # purposes. This endpoint is authenticated, whereas the one in PBAFileDownload can not be (at # the present time). In the future, consider whether or not they should be merged, or if they @@ -68,6 +72,8 @@ class FileUpload(AbstractResource): logger.error(str(err)) return make_response({"error": str(err)}, 404) + # NOTE: Consider putting most of this functionality into a service when this is transformed into + # a payload plugin. @jwt_required def post(self, target_os): """ @@ -85,6 +91,7 @@ class FileUpload(AbstractResource): try: self._update_config(target_os, safe_filename) except Exception as err: + # Roll back the entire transaction if part of it failed. self._file_storage.delete_file(safe_filename) raise err @@ -124,6 +131,7 @@ class FileUpload(AbstractResource): try: self._file_storage_service.delete_file(filename) except Exception as err: + # Roll back the entire transaction if part of it failed. self._agent_configuration_repository.store_configuration(original_agent_configuration) raise err From 9c7bf6c8b546c04eb198a3a09bf3f98f693d6285 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 10:27:47 -0400 Subject: [PATCH 4/6] Island: Use HTTPStatus Enum in PBA FileUpload resource --- monkey/monkey_island/cc/resources/pba_file_upload.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index eee10a3d9..61de8e1ec 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -70,7 +70,7 @@ class FileUpload(AbstractResource): return send_file(file, mimetype="application/octet-stream") except repository.FileNotFoundError as err: logger.error(str(err)) - return make_response({"error": str(err)}, 404) + return make_response({"error": str(err)}, HTTPStatus.NOT_FOUND) # NOTE: Consider putting most of this functionality into a service when this is transformed into # a payload plugin. @@ -96,7 +96,7 @@ class FileUpload(AbstractResource): raise err # API Spec: HTTP status code should be 201 - response = Response(response=safe_filename, status=200, mimetype="text/plain") + response = Response(response=safe_filename, status=HTTPStatus.OK, mimetype="text/plain") return response def _update_config(self, target_os: str, safe_filename: str): @@ -136,7 +136,7 @@ class FileUpload(AbstractResource): raise err # API Spec: HTTP status code should be 204 - return make_response({}, 200) + return make_response({}, HTTPStatus.OK) @staticmethod def _is_target_os_supported(target_os: str) -> bool: From e48e2cb9af34077783755277c450599da6c9e1ad Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 11:38:52 -0400 Subject: [PATCH 5/6] UT: Add typehints to test_pba_file_upload.py --- .../cc/resources/test_pba_file_upload.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 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 423dd262c..a1f39aaf6 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,13 @@ +from typing import Callable + import pytest +from flask.testing import FlaskClient from tests.common import StubDIContainer from tests.monkey_island import InMemoryAgentConfigurationRepository, SingleFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ +from common import DIContainer from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload @@ -25,17 +29,21 @@ Content-Type: text/x-python @pytest.fixture -def file_repository(): +def file_repository() -> IFileRepository: return SingleFileRepository() @pytest.fixture -def agent_configuration_repository(): +def agent_configuration_repository() -> IAgentConfigurationRepository: return InMemoryAgentConfigurationRepository() @pytest.fixture -def flask_client(build_flask_client, file_repository, agent_configuration_repository): +def flask_client( + build_flask_client: Callable[[DIContainer], FlaskClient], + file_repository: IFileRepository, + agent_configuration_repository: IAgentConfigurationRepository, +) -> FlaskClient: container = StubDIContainer() container.register_instance(IFileRepository, file_repository) container.register_instance(IAgentConfigurationRepository, agent_configuration_repository) @@ -45,7 +53,7 @@ def flask_client(build_flask_client, file_repository, agent_configuration_reposi @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post(flask_client, pba_os): +def test_pba_file_upload_post(flask_client: FlaskClient, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( url, @@ -56,7 +64,7 @@ def test_pba_file_upload_post(flask_client, pba_os): assert resp.status_code == 200 -def test_pba_file_upload_post__invalid(flask_client): +def test_pba_file_upload_post__invalid(flask_client: FlaskClient): url = get_url_for_resource(FileUpload, target_os="bogus") resp = flask_client.post( url, @@ -68,7 +76,9 @@ def test_pba_file_upload_post__invalid(flask_client): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_post__internal_server_error(flask_client, pba_os, file_repository): +def test_pba_file_upload_post__internal_server_error( + flask_client: FlaskClient, pba_os: str, file_repository: IFileRepository +): file_repository.save_file = lambda x, y: raise_(Exception()) url = get_url_for_resource(FileUpload, target_os=pba_os) @@ -82,14 +92,14 @@ def test_pba_file_upload_post__internal_server_error(flask_client, pba_os, file_ @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_get__file_not_found(flask_client, pba_os): +def test_pba_file_upload_get__file_not_found(flask_client: FlaskClient, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") resp = flask_client.get(url) assert resp.status_code == 404 @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_file_download_endpoint_500(open_error_flask_client, pba_os): +def test_file_download_endpoint_500(open_error_flask_client, pba_os: str): url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") resp = open_error_flask_client.get(url) @@ -98,7 +108,7 @@ def test_file_download_endpoint_500(open_error_flask_client, pba_os): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) -def test_pba_file_upload_endpoint(flask_client, pba_os): +def test_pba_file_upload_endpoint(flask_client: FlaskClient, pba_os: str): url_with_os = get_url_for_resource(FileUpload, target_os=pba_os) resp_post = flask_client.post( @@ -123,7 +133,7 @@ def test_pba_file_upload_endpoint(flask_client, pba_os): assert resp_get_del.status_code == 404 -def test_pba_file_upload_endpoint__invalid(flask_client): +def test_pba_file_upload_endpoint__invalid(flask_client: FlaskClient): url_with_os = get_url_for_resource(FileUpload, target_os="bogus") resp_post = flask_client.post( From a979a372fc7053fd156a60e3c7889015de654657 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 14 Jul 2022 06:55:25 -0400 Subject: [PATCH 6/6] Island: Rename _is_target_os_supported -> _target_os_is_unsupported --- monkey/monkey_island/cc/resources/pba_file_upload.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 61de8e1ec..8a67970c5 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -53,7 +53,7 @@ class FileUpload(AbstractResource): :param target_os: Indicates which file to send, linux or windows :return: Returns file contents """ - if self._is_target_os_supported(target_os): + if self._target_os_is_unsupported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") agent_configuration = self._agent_configuration_repository.get_configuration() @@ -81,7 +81,7 @@ class FileUpload(AbstractResource): :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_target_os_supported(target_os): + if self._target_os_is_unsupported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") file_storage = next(request.files.values()) # For now, assume there's only one file @@ -117,7 +117,7 @@ class FileUpload(AbstractResource): :param target_os: Type indicates which file was deleted, linux of windows :return: Empty response """ - if self._is_target_os_supported(target_os): + if self._target_os_is_unsupported(target_os): return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") original_agent_configuration = self._agent_configuration_repository.get_configuration() @@ -139,5 +139,5 @@ class FileUpload(AbstractResource): return make_response({}, HTTPStatus.OK) @staticmethod - def _is_target_os_supported(target_os: str) -> bool: + def _target_os_is_unsupported(target_os: str) -> bool: return target_os not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE}