Merge pull request #1320 from guardicore/1241/ransomware-quickstart-get-mode

Ransomware quickstart - Add get method for island mode
This commit is contained in:
VakarisZ 2021-07-14 12:25:53 +03:00 committed by GitHub
commit bf517bf566
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 24 additions and 9 deletions

View File

@ -5,7 +5,7 @@ import flask_restful
from flask import make_response, request from flask import make_response, request
from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.resources.auth.auth import jwt_required
from monkey_island.cc.services.mode.island_mode_service import set_mode from monkey_island.cc.services.mode.island_mode_service import get_mode, set_mode
from monkey_island.cc.services.mode.mode_enum import IslandModeEnum from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -26,3 +26,11 @@ class IslandMode(flask_restful.Resource):
return make_response({}, 400) return make_response({}, 400)
except ValueError: except ValueError:
return make_response({}, 422) return make_response({}, 422)
@jwt_required
def get(self):
try:
island_mode = get_mode()
return make_response({"mode": island_mode}, 200)
except IndexError:
return make_response({}, 422)

View File

@ -6,3 +6,8 @@ def set_mode(mode: IslandModeEnum):
island_mode_model = IslandMode() island_mode_model = IslandMode()
island_mode_model.mode = mode.value island_mode_model.mode = mode.value
island_mode_model.save() island_mode_model.save()
def get_mode():
mode = IslandMode.objects[0].mode
return mode

View File

@ -42,15 +42,17 @@ def test_island_mode_post__internal_server_error(monkeypatch, flask_client):
assert resp.status_code == 500 assert resp.status_code == 500
def test_island_mode_post__set_model(flask_client, uses_database): @pytest.mark.parametrize("mode", ["ransomware", "advanced"])
flask_client.post( def test_island_mode_endpoint(flask_client, uses_database, mode):
"/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True flask_client.post("/api/island-mode", data=json.dumps({"mode": mode}), follow_redirects=True)
) resp = flask_client.get("/api/island-mode", follow_redirects=True)
assert IslandMode.objects[0].mode == "ransomware" assert resp.status_code == 200
assert json.loads(resp.data)["mode"] == mode
def test_island_mode_post__set_invalid_model(flask_client, uses_database): def test_island_mode_endpoint__invalid_mode(flask_client, uses_database):
flask_client.post( flask_client.post(
"/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True "/api/island-mode", data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True
) )
assert len(IslandMode.objects) == 0 resp = flask_client.get("/api/island-mode", follow_redirects=True)
assert resp.status_code == 422