From 2fd38061b26aa2321add9f6a87497dc4b97dca6c Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 10 Sep 2021 14:01:58 +0200 Subject: [PATCH] UT: Add unit tests for pba_upload --- .../cc/resources/pba_file_upload.py | 18 +++ .../cc/resources/test_pba_file_upload.py | 120 ++++++++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 23554ab06..877aed437 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -27,6 +27,9 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file to send, linux or windows :return: Returns file contents """ + if self.check_file_type(file_type): + return Response(status=422, mimetype="text/plain") + # Verify that file_name is indeed a file from config if file_type == LINUX_PBA_TYPE: filename = ConfigService.get_config_value(copy.deepcopy(PBA_LINUX_FILENAME_PATH)) @@ -41,6 +44,9 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file was received, linux or windows :return: Returns flask response object with uploaded file's filename """ + if self.check_file_type(file_type): + return Response(status=422, mimetype="text/plain") + filename = FileUpload.upload_pba_file( request.files["filepond"], (file_type == LINUX_PBA_TYPE) ) @@ -74,6 +80,9 @@ class FileUpload(flask_restful.Resource): :param file_type: Type indicates which file was deleted, linux of windows :return: Empty response """ + if self.check_file_type(file_type): + return Response(status=422, mimetype="text/plain") + filename_path = ( PBA_LINUX_FILENAME_PATH if file_type == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH ) @@ -83,3 +92,12 @@ class FileUpload(flask_restful.Resource): ConfigService.set_config_value(filename_path, "") return {} + + @staticmethod + def check_file_type(file_type): + """ + Check if the file type is not supported + :param file_type: Type indicates which file was received, linux or windows + :return: Boolean + """ + return file_type not in {LINUX_PBA_TYPE, WINDOWS_PBA_TYPE} 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 new file mode 100644 index 000000000..085d555ea --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -0,0 +1,120 @@ +import pytest +from tests.utils import raise_ + +from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE +from monkey_island.cc.services.post_breach_files import PostBreachFilesService + +TEST_FILE = b"""-----------------------------1 +Content-Disposition: form-data; name="filepond" + +{} +-----------------------------1 +Content-Disposition: form-data; name="filepond"; filename="test.py" +Content-Type: text/x-python + +m0nk3y +-----------------------------1--""" + + +@pytest.fixture(autouse=True) +def custom_pba_directory(tmpdir): + PostBreachFilesService.initialize(tmpdir) + + +@pytest.fixture +def fake_set_config_value(monkeypatch): + monkeypatch.setattr( + "monkey_island.cc.services.config.ConfigService.set_config_value", lambda _, __: None + ) + + +@pytest.fixture +def fake_get_config_value(monkeypatch): + monkeypatch.setattr( + "monkey_island.cc.services.config.ConfigService.get_config_value", lambda _: "test.py" + ) + + +@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) +def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, fake_set_config_value): + resp = flask_client.post( + f"/api/fileUpload/{pba_os}", + data=TEST_FILE, + content_type="multipart/form-data; " "boundary=---------------------------" "1", + follow_redirects=True, + ) + assert resp.status_code == 200 + + +def test_pba_file_upload_post__invalid(flask_client, monkeypatch, fake_set_config_value): + resp = flask_client.post( + "/api/fileUpload/bogus", + data=TEST_FILE, + content_type="multipart/form-data; " "boundary=---------------------------" "1", + follow_redirects=True, + ) + assert resp.status_code == 422 + + +@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) +def test_pba_file_upload_post__internal_server_error( + flask_client, pba_os, monkeypatch, fake_set_config_value +): + monkeypatch.setattr( + "monkey_island.cc.resources.pba_file_upload.FileUpload.upload_pba_file", + lambda x, y: raise_(Exception()), + ) + + resp = flask_client.post( + f"/api/fileUpload/{pba_os}", + data=TEST_FILE, + content_type="multipart/form-data; boundary=---------------------------1", + follow_redirects=True, + ) + assert resp.status_code == 500 + + +@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) +def test_pba_file_upload_get__file_not_found( + flask_client, pba_os, monkeypatch, fake_get_config_value +): + resp = flask_client.get(f"/api/fileUpload/{pba_os}?load=bogus_mogus.py") + assert resp.status_code == 404 + + +@pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) +def test_pba_file_upload_endpoint( + flask_client, pba_os, monkeypatch, fake_get_config_value, fake_set_config_value +): + resp_post = flask_client.post( + f"/api/fileUpload/{pba_os}", + data=TEST_FILE, + content_type="multipart/form-data; " "boundary=---------------------------" "1", + follow_redirects=True, + ) + resp_get = flask_client.get(f"/api/fileUpload/{pba_os}?load=test.py") + resp_delete = flask_client.delete( + f"/api/fileUpload/{pba_os}", data="test.py", content_type="text/plain;" + ) + assert resp_post.status_code == 200 + assert resp_get.status_code == 200 + assert resp_get.data.decode() == "m0nk3y" + assert resp_delete.status_code == 200 + + +def test_pba_file_upload_endpoint__invalid( + flask_client, monkeypatch, fake_set_config_value, fake_get_config_value +): + resp_post = flask_client.post( + "/api/fileUpload/bogus", + data=TEST_FILE, + content_type="multipart/form-data; " "boundary=---------------------------" "1", + follow_redirects=True, + ) + resp_get = flask_client.get("/api/fileUpload/bogus?load=test.py") + resp_delete = flask_client.delete( + "/api/fileUpload/bogus", 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