From f8b6277a88d4d5b3083e126d8408a2003c6b9240 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 12 Jul 2021 13:32:14 +0300 Subject: [PATCH 01/12] Island: Add an endpoint for setting the island's mode. Also, add an enum of all the modes --- monkey/monkey_island/cc/app.py | 3 +++ .../monkey_island/cc/resources/island_mode.py | 20 +++++++++++++++++++ .../cc/services/mode/__init__.py | 0 .../cc/services/mode/island_mode_service.py | 5 +++++ .../cc/services/mode/mode_enum.py | 6 ++++++ 5 files changed, 34 insertions(+) create mode 100644 monkey/monkey_island/cc/resources/island_mode.py create mode 100644 monkey/monkey_island/cc/services/mode/__init__.py create mode 100644 monkey/monkey_island/cc/services/mode/island_mode_service.py create mode 100644 monkey/monkey_island/cc/services/mode/mode_enum.py diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 8800d382a..817a43333 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -26,6 +26,7 @@ from monkey_island.cc.resources.edge import Edge from monkey_island.cc.resources.environment import Environment from monkey_island.cc.resources.island_configuration import IslandConfiguration from monkey_island.cc.resources.island_logs import IslandLog +from monkey_island.cc.resources.island_mode import IslandMode from monkey_island.cc.resources.local_run import LocalRun from monkey_island.cc.resources.log import Log from monkey_island.cc.resources.monkey import Monkey @@ -132,6 +133,8 @@ def init_api_resources(api): api.add_resource( Telemetry, "/api/telemetry", "/api/telemetry/", "/api/telemetry/" ) + + api.add_resource(IslandMode, "/api/island-mode") api.add_resource(MonkeyConfiguration, "/api/configuration", "/api/configuration/") api.add_resource(IslandConfiguration, "/api/configuration/island", "/api/configuration/island/") api.add_resource(ConfigurationExport, "/api/configuration/export") diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py new file mode 100644 index 000000000..d0a109564 --- /dev/null +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -0,0 +1,20 @@ +import json + +import flask_restful +from flask import make_response, request + +from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.services.mode import island_mode_service +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +class IslandMode(flask_restful.Resource): + @jwt_required + def post(self): + body = json.loads(request.data) + mode_str = body.get("mode") + mode = IslandModeEnum(mode_str) + island_mode_service.set_mode(mode) + + # TODO return status + return make_response({}) diff --git a/monkey/monkey_island/cc/services/mode/__init__.py b/monkey/monkey_island/cc/services/mode/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py new file mode 100644 index 000000000..86f3e6d09 --- /dev/null +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -0,0 +1,5 @@ +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +def set_mode(mode: IslandModeEnum): + pass diff --git a/monkey/monkey_island/cc/services/mode/mode_enum.py b/monkey/monkey_island/cc/services/mode/mode_enum.py new file mode 100644 index 000000000..fce46db97 --- /dev/null +++ b/monkey/monkey_island/cc/services/mode/mode_enum.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class IslandModeEnum(Enum): + RANSOMWARE = "ransomware" + ADVANCED = "advanced" From 2778b69dfb4ef6a9f0c259a4c02c1fed8331f5af Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 12 Jul 2021 15:59:36 +0300 Subject: [PATCH 02/12] Island: Add unit test infrastructure for testing resources --- .../monkey_island/cc/resources/conftest.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py new file mode 100644 index 000000000..0e82fe163 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -0,0 +1,32 @@ +import flask_jwt_extended +import flask_restful +import pytest +from flask import Flask + +import monkey_island.cc.app +import monkey_island.cc.resources.auth.auth +import monkey_island.cc.resources.island_mode +from monkey_island.cc.services.representations import output_json + + +# We can't scope to module, because monkeypatch is a function scoped decorator. +# Potential solutions: https://github.com/pytest-dev/pytest/issues/363#issuecomment-406536200 or +# https://stackoverflow.com/questions/53963822/python-monkeypatch-setattr-with-pytest-fixture-at-module-scope +@pytest.fixture(scope="function") +def flask_client(monkeypatch): + monkeypatch.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) + + with mock_init_app().test_client() as client: + yield client + + +def mock_init_app(): + app = Flask(__name__) + + api = flask_restful.Api(app) + api.representations = {"application/json": output_json} + + monkey_island.cc.app.init_app_url_rules(app) + monkey_island.cc.app.init_api_resources(api) + + return app From 3bde6f013a4e4e2e6666414ae2e5e828f05437e1 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 12 Jul 2021 16:00:23 +0300 Subject: [PATCH 03/12] Island: Add a couple of island mode resource unit tests --- .../monkey_island/cc/resources/test_island_mode.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py new file mode 100644 index 000000000..01cb17aea --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -0,0 +1,13 @@ +import json + +import pytest + + +def test_island_mode_post(flask_client): + resp = flask_client.post('/api/island-mode', data=json.dumps({"mode": "ransomware"}), follow_redirects=True) + assert resp.status_code == 200 + + +def test_island_mode_post__invalid_mode(flask_client): + with pytest.raises(TypeError): + flask_client.post('/api/island-mode', data=json.dumps({"mode": "bogus mode"}), follow_redirects=True) From f9ed53a5272de4e5c50b5d0423b411288c993285 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 13 Jul 2021 10:30:12 +0200 Subject: [PATCH 04/12] Island: Add UT tests for island mode model --- .../cc/models/island_mode_model.py | 5 ++++ .../monkey_island/cc/resources/island_mode.py | 16 +++++++---- .../cc/services/mode/island_mode_service.py | 6 +++- .../monkey_island/cc/resources/conftest.py | 13 +++++++++ .../cc/resources/test_island_mode.py | 28 +++++++++++++++++-- 5 files changed, 58 insertions(+), 10 deletions(-) create mode 100644 monkey/monkey_island/cc/models/island_mode_model.py diff --git a/monkey/monkey_island/cc/models/island_mode_model.py b/monkey/monkey_island/cc/models/island_mode_model.py new file mode 100644 index 000000000..dec93e501 --- /dev/null +++ b/monkey/monkey_island/cc/models/island_mode_model.py @@ -0,0 +1,5 @@ +from mongoengine import Document, StringField + + +class IslandMode(Document): + mode = StringField() diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index d0a109564..cf8a3679e 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -1,20 +1,24 @@ import json +import logging import flask_restful from flask import make_response, request from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.services.mode import island_mode_service +from monkey_island.cc.services.mode.island_mode_service import set_mode from monkey_island.cc.services.mode.mode_enum import IslandModeEnum +logger = logging.getLogger(__name__) + class IslandMode(flask_restful.Resource): @jwt_required def post(self): body = json.loads(request.data) mode_str = body.get("mode") - mode = IslandModeEnum(mode_str) - island_mode_service.set_mode(mode) - - # TODO return status - return make_response({}) + try: + mode = IslandModeEnum(mode_str) + mode_value = set_mode(mode) + return make_response({"status": "MODE_FOUND", "mode": mode_value}, 200) + except ValueError: + return make_response({"status": "MODE_NOT_FOUND"}, 404) diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index 86f3e6d09..75041d2c1 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -1,5 +1,9 @@ +from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.services.mode.mode_enum import IslandModeEnum def set_mode(mode: IslandModeEnum): - pass + island_mode_model = IslandMode() + island_mode_model.mode = mode.value + island_mode_model.save() + return mode.value diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index 0e82fe163..2421e00aa 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,11 +1,13 @@ import flask_jwt_extended import flask_restful +import mongoengine import pytest from flask import Flask import monkey_island.cc.app import monkey_island.cc.resources.auth.auth import monkey_island.cc.resources.island_mode +from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.services.representations import output_json @@ -30,3 +32,14 @@ def mock_init_app(): monkey_island.cc.app.init_api_resources(api) return app + + +@pytest.fixture(scope="module", autouse=True) +def fake_mongo(): + mongoengine.disconnect() + mongoengine.connect("mongoenginetest", host="mongomock://localhost") + + +@pytest.fixture(scope="function") +def uses_database(): + IslandMode.objects().delete() diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 01cb17aea..13ded4799 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -2,12 +2,34 @@ import json import pytest +from monkey_island.cc.models.island_mode_model import IslandMode + def test_island_mode_post(flask_client): - resp = flask_client.post('/api/island-mode', data=json.dumps({"mode": "ransomware"}), follow_redirects=True) + resp = flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) assert resp.status_code == 200 def test_island_mode_post__invalid_mode(flask_client): - with pytest.raises(TypeError): - flask_client.post('/api/island-mode', data=json.dumps({"mode": "bogus mode"}), follow_redirects=True) + resp = flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + ) + assert resp.status_code == 404 + + +@pytest.mark.usefixtures("uses_database") +def test_island_mode_post__set_model(flask_client): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) + assert IslandMode.objects[0].mode == "ransomware" + + +@pytest.mark.usefixtures("uses_database") +def test_island_mode_post__set_invalid_model(flask_client): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + ) + assert len(IslandMode.objects) == 0 From 81a8ccf6731650d03d501b2cc7d280935c7f3a51 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 13 Jul 2021 11:17:35 +0200 Subject: [PATCH 05/12] Island: Return empty post status for island mode --- monkey/monkey_island/cc/resources/island_mode.py | 6 +++--- .../monkey_island/cc/services/mode/island_mode_service.py | 1 - vulture_allowlist.py | 1 + 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index cf8a3679e..889b25a26 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -18,7 +18,7 @@ class IslandMode(flask_restful.Resource): mode_str = body.get("mode") try: mode = IslandModeEnum(mode_str) - mode_value = set_mode(mode) - return make_response({"status": "MODE_FOUND", "mode": mode_value}, 200) + set_mode(mode) + return make_response({}, 200) except ValueError: - return make_response({"status": "MODE_NOT_FOUND"}, 404) + return make_response({}, 404) diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index 75041d2c1..05195e301 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -6,4 +6,3 @@ def set_mode(mode: IslandModeEnum): island_mode_model = IslandMode() island_mode_model.mode = mode.value island_mode_model.save() - return mode.value diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 618fabaa6..5a430dc6c 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -171,6 +171,7 @@ MONKEY_LINUX_RUNNING # unused variable (monkey/monkey_island/cc/services/utils/ import_status # monkey_island\cc\resources\configuration_import.py:19 config_schema # monkey_island\cc\resources\configuration_import.py:25 exception_stream # unused attribute (monkey_island/cc/server_setup.py:104) +ADVANCED # unused attribute (monkey/monkey_island/cc/services/mode/mode_enum.py:6:) # these are not needed for it to work, but may be useful extra information to understand what's going on WINDOWS_PBA_TYPE # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23) From 9310463f44886a72c542452c6f10214177892ed0 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 13 Jul 2021 12:27:00 +0200 Subject: [PATCH 06/12] UT: Refactor island mode test for set model --- .../monkey_island/cc/resources/conftest.py | 13 ------------- .../monkey_island/cc/resources/test_island_mode.py | 11 +++++++---- 2 files changed, 7 insertions(+), 17 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index 2421e00aa..0e82fe163 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,13 +1,11 @@ import flask_jwt_extended import flask_restful -import mongoengine import pytest from flask import Flask import monkey_island.cc.app import monkey_island.cc.resources.auth.auth import monkey_island.cc.resources.island_mode -from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.services.representations import output_json @@ -32,14 +30,3 @@ def mock_init_app(): monkey_island.cc.app.init_api_resources(api) return app - - -@pytest.fixture(scope="module", autouse=True) -def fake_mongo(): - mongoengine.disconnect() - mongoengine.connect("mongoenginetest", host="mongomock://localhost") - - -@pytest.fixture(scope="function") -def uses_database(): - IslandMode.objects().delete() diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 13ded4799..5c4362c64 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -5,6 +5,11 @@ import pytest from monkey_island.cc.models.island_mode_model import IslandMode +@pytest.fixture(scope="function") +def uses_database(): + IslandMode.objects().delete() + + def test_island_mode_post(flask_client): resp = flask_client.post( "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True @@ -19,16 +24,14 @@ def test_island_mode_post__invalid_mode(flask_client): assert resp.status_code == 404 -@pytest.mark.usefixtures("uses_database") -def test_island_mode_post__set_model(flask_client): +def test_island_mode_post__set_model(flask_client, uses_database): flask_client.post( "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True ) assert IslandMode.objects[0].mode == "ransomware" -@pytest.mark.usefixtures("uses_database") -def test_island_mode_post__set_invalid_model(flask_client): +def test_island_mode_post__set_invalid_model(flask_client, uses_database): flask_client.post( "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True ) From c56ca37bc0c2ec8990f5d2a4bdada8a43a2a133e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 10:09:53 -0400 Subject: [PATCH 07/12] Island: Respond with 422 instead of 404 from POST /api/island-mode --- monkey/monkey_island/cc/resources/island_mode.py | 2 +- .../unit_tests/monkey_island/cc/resources/test_island_mode.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 889b25a26..dd51d9fdb 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -21,4 +21,4 @@ class IslandMode(flask_restful.Resource): set_mode(mode) return make_response({}, 200) except ValueError: - return make_response({}, 404) + return make_response({}, 422) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 5c4362c64..91c6d0435 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -21,7 +21,7 @@ def test_island_mode_post__invalid_mode(flask_client): resp = flask_client.post( "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True ) - assert resp.status_code == 404 + assert resp.status_code == 422 def test_island_mode_post__set_model(flask_client, uses_database): From acdfeb858f3246b156ba21addb96d6eb802bf9fa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 10:30:38 -0400 Subject: [PATCH 08/12] Tests: Move raise_() to a reusable location --- .../monkey_island/cc/services/test_post_breach_files.py | 5 +---- monkey/tests/utils.py | 4 ++++ 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py index 5a2ddaa17..90a649a39 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py @@ -2,15 +2,12 @@ import os import pytest from tests.monkey_island.utils import assert_windows_permissions +from tests.utils import raise_ from monkey_island.cc.server_utils.file_utils import is_windows_os from monkey_island.cc.services.post_breach_files import PostBreachFilesService -def raise_(ex): - raise ex - - @pytest.fixture(autouse=True) def custom_pba_directory(tmpdir): PostBreachFilesService.initialize(tmpdir) diff --git a/monkey/tests/utils.py b/monkey/tests/utils.py index 2be032aad..8aea2d007 100644 --- a/monkey/tests/utils.py +++ b/monkey/tests/utils.py @@ -18,3 +18,7 @@ def hash_file(filepath: Path): sha256.update(block) return sha256.hexdigest() + + +def raise_(ex): + raise ex From 7549e64b413387ebec3400b4ea2589580c468ec7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 10:46:47 -0400 Subject: [PATCH 09/12] Island: Return 500 from POST /api/island-mode if unexpected exception --- monkey/monkey_island/cc/resources/island_mode.py | 2 ++ .../monkey_island/cc/resources/test_island_mode.py | 11 +++++++++++ 2 files changed, 13 insertions(+) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index dd51d9fdb..b69b399c0 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -22,3 +22,5 @@ class IslandMode(flask_restful.Resource): return make_response({}, 200) except ValueError: return make_response({}, 422) + except Exception: + return make_response({}, 500) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 91c6d0435..df4e9c25a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -1,8 +1,10 @@ import json import pytest +from tests.utils import raise_ from monkey_island.cc.models.island_mode_model import IslandMode +from monkey_island.cc.resources import island_mode as island_mode_resource @pytest.fixture(scope="function") @@ -24,6 +26,15 @@ def test_island_mode_post__invalid_mode(flask_client): assert resp.status_code == 422 +def test_island_mode_post__internal_server_error(monkeypatch, flask_client): + monkeypatch.setattr(island_mode_resource, "set_mode", lambda x: raise_(Exception())) + + resp = flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) + assert resp.status_code == 500 + + def test_island_mode_post__set_model(flask_client, uses_database): flask_client.post( "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True From 26d3782a666db22095f01182f469a350d2961a73 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 10:49:15 -0400 Subject: [PATCH 10/12] Island: Test both "ransomware" and "advanced" modes --- .../monkey_island/cc/resources/test_island_mode.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index df4e9c25a..9bd4d8dd3 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -12,9 +12,10 @@ def uses_database(): IslandMode.objects().delete() -def test_island_mode_post(flask_client): +@pytest.mark.parametrize("mode", ["ransomware", "advanced"]) +def test_island_mode_post(flask_client, mode): resp = flask_client.post( - "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + "/api/island-mode", data=json.dumps({"mode": mode}), follow_redirects=True ) assert resp.status_code == 200 From a0fb6fa2b69378b4d2e011d9ef29301036b59e1c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 10:58:08 -0400 Subject: [PATCH 11/12] Island: Return 400 from POST /api/island_mode on invalid JSON --- monkey/monkey_island/cc/resources/island_mode.py | 8 ++++++-- .../monkey_island/cc/resources/test_island_mode.py | 6 ++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index b69b399c0..5a83fb46a 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -14,12 +14,16 @@ logger = logging.getLogger(__name__) class IslandMode(flask_restful.Resource): @jwt_required def post(self): - body = json.loads(request.data) - mode_str = body.get("mode") try: + body = json.loads(request.data) + mode_str = body.get("mode") + mode = IslandModeEnum(mode_str) set_mode(mode) + return make_response({}, 200) + except (AttributeError, json.decoder.JSONDecodeError): + return make_response({}, 400) except ValueError: return make_response({}, 422) except Exception: diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 9bd4d8dd3..b8d6a84c2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -27,6 +27,12 @@ def test_island_mode_post__invalid_mode(flask_client): assert resp.status_code == 422 +@pytest.mark.parametrize("invalid_json", ["42", "{test"]) +def test_island_mode_post__invalid_json(flask_client, invalid_json): + resp = flask_client.post("/api/island-mode", data="{test", follow_redirects=True) + assert resp.status_code == 400 + + def test_island_mode_post__internal_server_error(monkeypatch, flask_client): monkeypatch.setattr(island_mode_resource, "set_mode", lambda x: raise_(Exception())) From 84a78a5048797db91556729000555cff5692fd78 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 13 Jul 2021 11:02:18 -0400 Subject: [PATCH 12/12] Island: Don't catch Exception in POST /api/island-mode Flask automatically traps exceptions, returns a 500, and logs a stack trace. Since Flask will automatically return a 500, we don't need to duplicate the functionality. Since it prints a stack trace, it provides more useful information than catching it did. --- monkey/monkey_island/cc/resources/island_mode.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 5a83fb46a..51df90e2d 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -26,5 +26,3 @@ class IslandMode(flask_restful.Resource): return make_response({}, 400) except ValueError: return make_response({}, 422) - except Exception: - return make_response({}, 500)