Island: Remove ConfigService from PBA FileUpload resource
This commit is contained in:
Mike Salvatore
parent
b4ced896b6
commit
2e7bcd54df
|
|
@ -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/<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
|
||||
|
||||
# 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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in New Issue