Merge pull request #1320 from guardicore/1241/ransomware-quickstart-get-mode
Ransomware quickstart - Add get method for island mode
This commit is contained in:
commit
bf517bf566
|
@ -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)
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue