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 570d3239c..d54d14f84 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 @@ -3,7 +3,9 @@ from typing import BinaryIO import pytest from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import get_url_for_resource +from monkey_island.cc.resources.pba_file_download import PBAFileDownload from monkey_island.cc.services import FileRetrievalError, IFileStorageService FILE_NAME = "test_file" @@ -40,7 +42,8 @@ def flask_client(build_flask_client): def test_file_download_endpoint(tmp_path, flask_client): - resp = flask_client.get(f"/api/pba/download/{FILE_NAME}") + download_url = get_url_for_resource(PBAFileDownload, filename=FILE_NAME) + resp = flask_client.get(download_url) assert resp.status_code == 200 assert next(resp.response) == FILE_CONTENTS @@ -48,7 +51,8 @@ def test_file_download_endpoint(tmp_path, flask_client): def test_file_download_endpoint_404(tmp_path, flask_client): nonexistant_file_name = "nonexistant_file" + download_url = get_url_for_resource(PBAFileDownload, filename=nonexistant_file_name) - resp = flask_client.get(f"/api/pba/download/{nonexistant_file_name}") + resp = flask_client.get(download_url) 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 9cbaa50d8..3ec9a631a 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 @@ -3,9 +3,10 @@ from typing import BinaryIO import pytest from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ -from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE +from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload from monkey_island.cc.services import FileRetrievalError, IFileStorageService TEST_FILE_CONTENTS = b"m0nk3y" @@ -74,8 +75,9 @@ 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, mock_set_config_value): + url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( - f"/api/file-upload/{pba_os}", + url, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, @@ -84,8 +86,9 @@ def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): + url = get_url_for_resource(FileUpload, target_os="bogus") resp = flask_client.post( - "/api/file-upload/bogus", + url, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, @@ -98,9 +101,10 @@ def test_pba_file_upload_post__internal_server_error( flask_client, pba_os, mock_set_config_value, file_storage_service ): file_storage_service.save_file = lambda x, y: raise_(Exception()) + url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( - f"/api/file-upload/{pba_os}", + url, data=TEST_FILE, content_type="multipart/form-data; boundary=---------------------------1", follow_redirects=True, @@ -110,7 +114,8 @@ 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): - resp = flask_client.get(f"/api/file-upload/{pba_os}?load=bogus_mogus.py") + url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") + resp = flask_client.get(url) assert resp.status_code == 404 @@ -118,23 +123,24 @@ def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_conf def test_pba_file_upload_endpoint( 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( - f"/api/file-upload/{pba_os}", + url_with_os, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, ) - resp_get = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") + url_with_filename = get_url_for_resource(FileUpload, target_os=pba_os, filename="test.py") + resp_get = flask_client.get(url_with_filename) assert resp_get.status_code == 200 assert resp_get.data == TEST_FILE_CONTENTS # Closing the response closes the file handle, else it can't be deleted resp_get.close() - resp_delete = flask_client.delete( - f"/api/file-upload/{pba_os}", data="test.py", content_type="text/plain;" - ) - resp_get_del = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") + 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 @@ -145,16 +151,20 @@ def test_pba_file_upload_endpoint( def test_pba_file_upload_endpoint__invalid( 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( - "/api/file-upload/bogus", + url_with_os, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, ) - resp_get = flask_client.get("/api/file-upload/bogus?load=test.py") - resp_delete = flask_client.delete( - "/api/file-upload/bogus", data="test.py", content_type="text/plain;" + + url_with_filename = get_url_for_resource( + FileUpload, target_os="bogus", filename="bobug_mogus.py" ) + resp_get = flask_client.get(url_with_filename) + resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;") assert resp_post.status_code == 422 assert resp_get.status_code == 422 assert resp_delete.status_code == 422 diff --git a/monkey/tests/unit_tests/monkey_island/conftest.py b/monkey/tests/unit_tests/monkey_island/conftest.py index 66780882c..03cc23259 100644 --- a/monkey/tests/unit_tests/monkey_island/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/conftest.py @@ -1,11 +1,14 @@ import os +import re from collections.abc import Callable +from typing import Set import flask_restful import pytest from flask import Flask import monkey_island +from monkey_island.cc.resources.i_resource import IResource from monkey_island.cc.services.representations import output_json @@ -37,3 +40,27 @@ def mock_flask_resource_manager(container): flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) return flask_resource_manager + + +def get_url_for_resource(resource: IResource, **kwargs): + chosen_url = None + for url in resource.urls: + if _get_url_keywords(url) == set(kwargs.keys()): + chosen_url = url + if not chosen_url: + raise Exception( + f"Resource {resource} doesn't contain a url that matches {kwargs} keywords." + ) + + for key, value in kwargs.items(): + reg_pattern = f"<.*:{key}>" + chosen_url = re.sub(pattern=reg_pattern, repl=value, string=chosen_url) + + return chosen_url + + +def _get_url_keywords(url: str) -> Set[str]: + # Match pattern , but only put "keyword" in a group + reg_pattern = "(?:<.*?:)(.*?)(?:>)" + reg_matches = re.finditer(reg_pattern, url) + return set([match.groups()[0] for match in reg_matches])