Merge pull request #1318 from guardicore/ransomware_quickstart_endpoint

Ransomware quickstart endpoint

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/<string:monkey_guid>"
     )
+
+    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")

monkey/monkey_island/cc/models/island_mode_model.py

@@ -0,0 +1,5 @@
+from mongoengine import Document, StringField
+
+
+class IslandMode(Document):
+    mode = StringField()

monkey/monkey_island/cc/resources/island_mode.py

@@ -0,0 +1,28 @@
+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.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):
+        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)

monkey/monkey_island/cc/services/mode/island_mode_service.py

@@ -0,0 +1,8 @@
+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):
+    island_mode_model = IslandMode()
+    island_mode_model.mode = mode.value
+    island_mode_model.save()

monkey/monkey_island/cc/services/mode/mode_enum.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class IslandModeEnum(Enum):
+    RANSOMWARE = "ransomware"
+    ADVANCED = "advanced"

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

monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py

@@ -0,0 +1,56 @@
+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")
+def uses_database():
+    IslandMode.objects().delete()
+
+
+@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": mode}), follow_redirects=True
+    )
+    assert resp.status_code == 200
+
+
+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 == 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()))
+
+    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
+    )
+    assert IslandMode.objects[0].mode == "ransomware"
+
+
+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
+    )
+    assert len(IslandMode.objects) == 0

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)

monkey/tests/utils.py

@@ -18,3 +18,7 @@ def hash_file(filepath: Path):
             sha256.update(block)
 
     return sha256.hexdigest()
+
+
+def raise_(ex):
+    raise ex

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)