Merge pull request #2087 from guardicore/1662-remove-configservice-from-pba-resource

1662 remove configservice from pba resource

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__)
 
@@ -18,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.
     """
@@ -30,8 +30,16 @@ class FileUpload(AbstractResource):
         "/api/file-upload/<string:target_os>?restore=<string:filename>",
     ]
 
-    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
+
+    # 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
@@ -45,14 +53,15 @@ 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()
         # 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)
@@ -61,8 +70,10 @@ 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.
     @jwt_required
     def post(self, target_os):
         """
@@ -70,22 +81,35 @@ 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
         safe_filename = sanitize_filename(file_storage.filename)
 
         self._file_storage_service.save_file(safe_filename, file_storage.stream)
-        ConfigService.set_config_value(
+        try:
-            (PBA_LINUX_FILENAME_PATH if target_os == LINUX_PBA_TYPE else PBA_WINDOWS_FILENAME_PATH),
+            self._update_config(target_os, safe_filename)
-            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
 
         # 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):
+        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):
         """
@@ -93,20 +117,27 @@ 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")
 
-        filename_path = (
+        original_agent_configuration = self._agent_configuration_repository.get_configuration()
-            PBA_LINUX_FILENAME_PATH if target_os == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH
+        self._update_config(target_os, "")
-        )
+
-        filename = ConfigService.get_config_value(filename_path)
+        if target_os == LINUX_PBA_TYPE:
-        if filename:
+            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:
+            # Roll back the entire transaction if part of it failed.
+            self._agent_configuration_repository.store_configuration(original_agent_configuration)
+            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:
+    def _target_os_is_unsupported(target_os: str) -> bool:
         return target_os not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE}

monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py

@@ -1,10 +1,14 @@
+from typing import Callable
+
 import pytest
+from flask.testing import FlaskClient
 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 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
 
 TEST_FILE_CONTENTS = b"m0nk3y"
@@ -25,35 +29,31 @@ Content-Type: text/x-python
 
 
 @pytest.fixture
-def mock_set_config_value(monkeypatch):
+def file_repository() -> IFileRepository:
-    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() -> IAgentConfigurationRepository:
+    return InMemoryAgentConfigurationRepository()
+
+
+@pytest.fixture
+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)
 
     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: FlaskClient, pba_os: str):
     url = get_url_for_resource(FileUpload, target_os=pba_os)
     resp = flask_client.post(
         url,
@@ -64,7 +64,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: FlaskClient):
     url = get_url_for_resource(FileUpload, target_os="bogus")
     resp = flask_client.post(
         url,
@@ -77,7 +77,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
+    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)
@@ -92,14 +92,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, mock_get_config_value):
+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)
@@ -108,9 +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(
+def test_pba_file_upload_endpoint(flask_client: FlaskClient, pba_os: str):
-    flask_client, pba_os, mock_get_config_value, mock_set_config_value
-):
 
     url_with_os = get_url_for_resource(FileUpload, target_os=pba_os)
     resp_post = flask_client.post(
@@ -129,16 +127,13 @@ 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
 
 
-def test_pba_file_upload_endpoint__invalid(
+def test_pba_file_upload_endpoint__invalid(flask_client: FlaskClient):
-    flask_client, mock_set_config_value, mock_get_config_value
-):
 
     url_with_os = get_url_for_resource(FileUpload, target_os="bogus")
     resp_post = flask_client.post(