diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 877aed437..68be20e9a 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -1,4 +1,5 @@ import copy +from http import HTTPStatus import flask_restful from flask import Response, request, send_from_directory @@ -27,8 +28,8 @@ 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") + if self.is_pba_file_type_supported(file_type): + return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") # Verify that file_name is indeed a file from config if file_type == LINUX_PBA_TYPE: @@ -44,8 +45,8 @@ 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") + if self.is_pba_file_type_supported(file_type): + return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") filename = FileUpload.upload_pba_file( request.files["filepond"], (file_type == LINUX_PBA_TYPE) @@ -80,8 +81,8 @@ 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") + if self.is_pba_file_type_supported(file_type): + return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY, mimetype="text/plain") filename_path = ( PBA_LINUX_FILENAME_PATH if file_type == "PBAlinux" else PBA_WINDOWS_FILENAME_PATH @@ -94,10 +95,5 @@ class FileUpload(flask_restful.Resource): 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 - """ + def is_pba_file_type_supported(file_type: str) -> bool: 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 index 085d555ea..c06f0c400 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 @@ -22,21 +22,21 @@ def custom_pba_directory(tmpdir): @pytest.fixture -def fake_set_config_value(monkeypatch): +def mock_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): +def mock_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): +def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, mock_set_config_value): resp = flask_client.post( f"/api/fileUpload/{pba_os}", data=TEST_FILE, @@ -46,7 +46,7 @@ def test_pba_file_upload_post(flask_client, pba_os, monkeypatch, fake_set_config assert resp.status_code == 200 -def test_pba_file_upload_post__invalid(flask_client, monkeypatch, fake_set_config_value): +def test_pba_file_upload_post__invalid(flask_client, monkeypatch, mock_set_config_value): resp = flask_client.post( "/api/fileUpload/bogus", data=TEST_FILE, @@ -58,7 +58,7 @@ def test_pba_file_upload_post__invalid(flask_client, monkeypatch, fake_set_confi @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 + flask_client, pba_os, monkeypatch, mock_set_config_value ): monkeypatch.setattr( "monkey_island.cc.resources.pba_file_upload.FileUpload.upload_pba_file", @@ -76,7 +76,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, monkeypatch, fake_get_config_value + flask_client, pba_os, monkeypatch, mock_get_config_value ): resp = flask_client.get(f"/api/fileUpload/{pba_os}?load=bogus_mogus.py") assert resp.status_code == 404 @@ -84,7 +84,7 @@ def test_pba_file_upload_get__file_not_found( @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 + flask_client, pba_os, monkeypatch, mock_get_config_value, mock_set_config_value ): resp_post = flask_client.post( f"/api/fileUpload/{pba_os}", @@ -96,14 +96,19 @@ def test_pba_file_upload_endpoint( resp_delete = flask_client.delete( f"/api/fileUpload/{pba_os}", data="test.py", content_type="text/plain;" ) + resp_get_del = flask_client.get(f"/api/fileUpload/{pba_os}?load=test.py") assert resp_post.status_code == 200 + assert resp_get.status_code == 200 assert resp_get.data.decode() == "m0nk3y" + assert resp_delete.status_code == 200 + assert resp_get_del.status_code == 404 + def test_pba_file_upload_endpoint__invalid( - flask_client, monkeypatch, fake_set_config_value, fake_get_config_value + flask_client, monkeypatch, mock_set_config_value, mock_get_config_value ): resp_post = flask_client.post( "/api/fileUpload/bogus",