From 86242a021d55199f37c18826b5f8c99173bf84b0 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 1 Aug 2022 18:45:00 +0200 Subject: [PATCH] Island: Return island mode only as a string --- .../monkey_island/cc/resources/island_mode.py | 7 +++--- .../cc/resources/test_island_mode.py | 22 ++++++------------- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 5bd695b3d..07d5d8095 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -2,7 +2,7 @@ import json import logging from http import HTTPStatus -from flask import make_response, request +from flask import request from monkey_island.cc.models import IslandMode as IslandModeEnum from monkey_island.cc.resources.AbstractResource import AbstractResource @@ -21,8 +21,7 @@ class IslandMode(AbstractResource): @jwt_required def put(self): try: - body = json.loads(request.data) - mode = IslandModeEnum(body.get("mode")) + mode = IslandModeEnum(json.loads(request.data)) self._island_mode_service.set_mode(mode) @@ -35,4 +34,4 @@ class IslandMode(AbstractResource): @jwt_required def get(self): island_mode = self._island_mode_service.get_mode() - return make_response({"mode": island_mode.value}, HTTPStatus.OK) + return island_mode.value, HTTPStatus.OK 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 642c5c936..97aa300e0 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 @@ -38,24 +38,18 @@ def flask_client(build_flask_client): ) def test_island_mode_post(flask_client, mode): resp = flask_client.put( - IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps(mode), follow_redirects=True ) assert resp.status_code == HTTPStatus.NO_CONTENT def test_island_mode_post__invalid_mode(flask_client): resp = flask_client.put( - IslandModeResource.urls[0], data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps("bogus mode"), follow_redirects=True ) assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY -@pytest.mark.parametrize("invalid_json", ["42", "{test"]) -def test_island_mode_post__invalid_json(flask_client, invalid_json): - resp = flask_client.put(IslandModeResource.urls[0], data="{test", follow_redirects=True) - assert resp.status_code == HTTPStatus.BAD_REQUEST - - def test_island_mode_post__internal_server_error(build_flask_client): mock_island_mode_service = MagicMock(spec=IslandModeService) mock_island_mode_service.set_mode = MagicMock(side_effect=RetrievalError) @@ -66,7 +60,7 @@ def test_island_mode_post__internal_server_error(build_flask_client): with build_flask_client(container) as flask_client: resp = flask_client.put( IslandModeResource.urls[0], - data=json.dumps({"mode": IslandMode.RANSOMWARE.value}), + data=json.dumps(IslandMode.RANSOMWARE.value), follow_redirects=True, ) @@ -75,18 +69,16 @@ def test_island_mode_post__internal_server_error(build_flask_client): @pytest.mark.parametrize("mode", [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value]) def test_island_mode_endpoint(flask_client, mode): - flask_client.put( - IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True - ) + flask_client.put(IslandModeResource.urls[0], data=json.dumps(mode), follow_redirects=True) resp = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp.status_code == HTTPStatus.OK - assert json.loads(resp.data)["mode"] == mode + assert json.loads(resp.data) == mode def test_island_mode_endpoint__invalid_mode(flask_client): resp_post = flask_client.put( - IslandModeResource.urls[0], data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps("bogus_mode"), follow_redirects=True ) resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp_post.status_code == HTTPStatus.UNPROCESSABLE_ENTITY - assert json.loads(resp_get.data)["mode"] == IslandMode.UNSET.value + assert json.loads(resp_get.data) == IslandMode.UNSET.value