p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

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

1662 remove configservice from pba resource
This commit is contained in:
Mike Salvatore 2022-07-14 07:20:39 -04:00 committed by GitHub
parent c5dfd9c3e6 a979a372fc
commit ac704471d3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 80 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
monkey/monkey_island/cc/resources/pba_file_upload.py
View File

@ -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(
(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:
# 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 = (
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:
# 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}

57
monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py
View File

@ -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):
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():
def file_repository() -> IFileRepository:
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(
flask_client, pba_os, mock_get_config_value, mock_set_config_value
):
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(
@ -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(
flask_client, mock_set_config_value, mock_get_config_value
):
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(