Merge branch '2157-roles-flask' into develop

Issue #2157
PR #3064

Author: mssalvatore

monkey/common/__init__.py

@@ -9,3 +9,4 @@
 from .agent_signals import AgentSignals
 from .agent_heartbeat import AgentHeartbeat
 from .hard_coded_manifests import HARD_CODED_EXPLOITER_MANIFESTS
+from .account_role import AccountRole

monkey/common/account_role.py

@@ -0,0 +1,11 @@
+from enum import Enum, auto
+
+
+class AccountRole(Enum):
+    """
+    An Enum representing  roles.
+    This Enum represents roles that an account can have.
+    """
+
+    ISLAND_INTERFACE = auto()
+    AGENT = auto()

monkey/monkey_island/cc/app.py

@@ -5,11 +5,11 @@
 from flask import Flask, Response, send_from_directory
 from flask.sessions import SecureCookieSessionInterface
 from flask_mongoengine import MongoEngine
-from flask_security import ConfirmRegisterForm, MongoEngineUserDatastore, Security
+from flask_security import ConfirmRegisterForm, MongoEngineUserDatastore, Security, UserDatastore
 from werkzeug.exceptions import NotFound
 from wtforms import StringField, ValidationError
 
-from common import DIContainer
+from common import AccountRole, DIContainer
 from monkey_island.cc.flask_utils import FlaskDIWrapper
 from monkey_island.cc.models import Role, User
 from monkey_island.cc.resources import (
@@ -91,16 +91,16 @@ def setup_authentication(app, data_dir):
 
     # The database object needs to be created after we configure the flask application
     db = MongoEngine(app)
-
     user_datastore = MongoEngineUserDatastore(db, User, Role)
 
+    _create_roles(user_datastore)
+
     # Only one user can be registered in the Island, so we need a custom validator
     def validate_no_user_exists_already(_, field):
         if user_datastore.find_user():
             raise ValidationError("A user already exists. Only a single user can be registered.")
 
     class CustomConfirmRegisterForm(ConfirmRegisterForm):
-
         # We don't use the email, but the field is required by ConfirmRegisterForm.
         # Email validators need to be overriden, otherwise an error about invalid email is raised.
         # Added custom validator to the email field because we have to override
@@ -109,6 +109,11 @@ class CustomConfirmRegisterForm(ConfirmRegisterForm):
             "Email", default="dummy@dummy.com", validators=[validate_no_user_exists_already]
         )
 
+        def to_dict(self, only_user):
+            registration_dict = super().to_dict(only_user)
+            registration_dict.update({"roles": [AccountRole.ISLAND_INTERFACE.name]})
+            return registration_dict
+
     app.security = Security(
         app,
         user_datastore,
@@ -121,6 +126,11 @@ class CustomConfirmRegisterForm(ConfirmRegisterForm):
     app.session_interface = disable_session_cookies()
 
 
+def _create_roles(user_datastore: UserDatastore):
+    user_datastore.find_or_create_role(name=AccountRole.ISLAND_INTERFACE.name)
+    user_datastore.find_or_create_role(name=AccountRole.AGENT.name)
+
+
 def init_app_config(app, mongo_url, data_dir: Path):
     app.config["MONGO_URI"] = mongo_url
     app.config["MONGODB_SETTINGS"] = [
@@ -207,7 +217,11 @@ def init_rpc_endpoints(api: FlaskDIWrapper):
     api.add_resource(TerminateAllAgents)
 
 
-def init_app(mongo_url: str, container: DIContainer, data_dir: Path):
+def init_app(
+    mongo_url: str,
+    container: DIContainer,
+    data_dir: Path,
+):
     """
     Simple docstirng for init_app
 

monkey/monkey_island/cc/models/role.py

@@ -6,5 +6,3 @@
 
 class Role(Document, RoleMixin):
     name = StringField(max_length=80, unique=True)
-    description = StringField(max_length=255)
-    permissions = StringField(max_length=255)

monkey/monkey_island/cc/resources/agent_logs.py

@@ -2,8 +2,9 @@
 from http import HTTPStatus
 
 from flask import request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from common.types import AgentID
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import IAgentLogRepository, UnknownRecordError
@@ -18,6 +19,7 @@ def __init__(self, agent_log_repository: IAgentLogRepository):
         self._agent_log_repository = agent_log_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self, agent_id: AgentID):
         try:
             log_contents = self._agent_log_repository.get_agent_log(agent_id)

monkey/monkey_island/cc/resources/agent_signals/terminate_all_agents.py

@@ -3,8 +3,9 @@
 from json import JSONDecodeError
 
 from flask import request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.models import TerminateAllAgents as TerminateAllAgentsObject
@@ -22,6 +23,7 @@ def __init__(
         self._island_event_queue = island_event_queue
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def post(self):
         try:
             terminate_all_agents = TerminateAllAgentsObject(**request.json)

monkey/monkey_island/cc/resources/agents.py

@@ -3,9 +3,9 @@
 from http import HTTPStatus
 
 from flask import make_response, request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
-from common import AgentRegistrationData
+from common import AccountRole, AgentRegistrationData
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import IAgentRepository
@@ -21,6 +21,7 @@ def __init__(self, island_event_queue: IIslandEventQueue, agent_repository: IAge
         self._agent_repository = agent_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         return self._agent_repository.get_agents(), HTTPStatus.OK
 

monkey/monkey_island/cc/resources/clear_simulation_data.py

@@ -1,8 +1,9 @@
 from http import HTTPStatus
 
 from flask import make_response
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.flask_utils import AbstractResource
 
@@ -14,6 +15,7 @@ def __init__(self, island_event_queue: IIslandEventQueue):
         self._island_event_queue = island_event_queue
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def post(self):
         """
         Clear all data collected during the simulation

monkey/monkey_island/cc/resources/exploitations/monkey_exploitation.py

@@ -1,7 +1,8 @@
 from dataclasses import asdict
 
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import (
     IAgentEventRepository,
@@ -27,6 +28,7 @@ def __init__(
         self._agent_plugin_repository = agent_plugin_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         monkey_exploitations = [
             asdict(exploitation)

monkey/monkey_island/cc/resources/island_log.py

@@ -1,8 +1,9 @@
 import logging
 from pathlib import Path
 
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from common.utils.file_utils import get_text_file_contents
 from monkey_island.cc.flask_utils import AbstractResource
 
@@ -16,6 +17,7 @@ def __init__(self, island_log_file_path: Path):
         self._island_log_file_path = island_log_file_path
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         try:
             return get_text_file_contents(self._island_log_file_path)

monkey/monkey_island/cc/resources/island_mode.py

@@ -3,8 +3,9 @@
 from http import HTTPStatus
 
 from flask import request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.models import IslandMode as IslandModeEnum
@@ -25,6 +26,7 @@ def __init__(
         self._simulation_repository = simulation_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def put(self):
         try:
             mode = IslandModeEnum(request.json)
@@ -36,6 +38,7 @@ def put(self):
             return {}, HTTPStatus.UNPROCESSABLE_ENTITY
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         island_mode = self._simulation_repository.get_mode()
         return island_mode.value, HTTPStatus.OK

monkey/monkey_island/cc/resources/local_run.py

@@ -1,8 +1,9 @@
 import json
 
 from flask import jsonify, make_response, request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
 
@@ -15,6 +16,7 @@ def __init__(self, local_monkey_run_service: LocalMonkeyRunService):
 
     # API Spec: This should be an RPC-style endpoint
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def post(self):
         body = json.loads(request.data)
         if body.get("action") == "run":

monkey/monkey_island/cc/resources/machines.py

@@ -1,7 +1,8 @@
 from http import HTTPStatus
 
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import IMachineRepository
 
@@ -13,5 +14,6 @@ def __init__(self, machine_repository: IMachineRepository):
         self._machine_repository = machine_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         return self._machine_repository.get_machines(), HTTPStatus.OK

monkey/monkey_island/cc/resources/nodes.py

@@ -1,7 +1,8 @@
 from http import HTTPStatus
 
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import INodeRepository
 
@@ -13,5 +14,6 @@ def __init__(self, node_repository: INodeRepository):
         self._node_repository = node_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         return self._node_repository.get_nodes(), HTTPStatus.OK

monkey/monkey_island/cc/resources/ransomware_report.py

@@ -1,6 +1,7 @@
 from flask import jsonify
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import (
     IAgentEventRepository,
@@ -24,6 +25,7 @@ def __init__(
         self._agent_plugin_repository = agent_plugin_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         return jsonify(
             {

monkey/monkey_island/cc/resources/remote_run.py

@@ -3,8 +3,9 @@
 
 from botocore.exceptions import ClientError, NoCredentialsError
 from flask import jsonify, make_response, request
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.services import AWSService
 from monkey_island.cc.services.aws import AWSCommandResults
@@ -29,6 +30,7 @@ def __init__(self, aws_service: AWSService):
         self._aws_service = aws_service
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         action = request.args.get("action")
         if action == "list_aws":
@@ -50,6 +52,7 @@ def get(self):
         return {}
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def post(self):
         body = json.loads(request.data)
         if body.get("type") == "aws":

monkey/monkey_island/cc/resources/report_generation_status.py

@@ -1,6 +1,7 @@
 from flask import jsonify
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.repositories import IAgentRepository
 from monkey_island.cc.services.infection_lifecycle import is_report_done
@@ -13,6 +14,7 @@ def __init__(self, agent_repository: IAgentRepository):
         self._agent_repository = agent_repository
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         return self.report_generation_status()
 

monkey/monkey_island/cc/resources/reset_agent_configuration.py

@@ -1,8 +1,9 @@
 from http import HTTPStatus
 
 from flask import make_response
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.flask_utils import AbstractResource
 
@@ -14,6 +15,7 @@ def __init__(self, island_event_queue: IIslandEventQueue):
         self._island_event_queue = island_event_queue
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def post(self):
         """
         Reset the agent configuration to its default values

monkey/monkey_island/cc/resources/security_report.py

@@ -1,5 +1,6 @@
-from flask_security import auth_token_required
+from flask_security import auth_token_required, roles_required
 
+from common import AccountRole
 from monkey_island.cc.flask_utils import AbstractResource
 from monkey_island.cc.services.reporting.report import ReportService
 
@@ -8,6 +9,7 @@ class SecurityReport(AbstractResource):
     urls = ["/api/report/security"]
 
     @auth_token_required
+    @roles_required(AccountRole.ISLAND_INTERFACE.name)
     def get(self):
         ReportService.update_report()
         return ReportService.get_report()