diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py index d3a36e6a2..aa2087913 100644 --- a/monkey/monkey_island/cc/resources/root.py +++ b/monkey/monkey_island/cc/resources/root.py @@ -19,6 +19,8 @@ class Root(flask_restful.Resource): if not action: return self.get_server_info() + elif action == "delete-agent-data": + return jwt_required(Database.reset_db)(reset_config=False) elif action == "reset": return jwt_required(Database.reset_db)() elif action == "is-up": diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index 7aeb1bfcf..46b5e0ffd 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -3,8 +3,10 @@ import logging from flask import jsonify from monkey_island.cc.database import mongo +from monkey_island.cc.models import Config from monkey_island.cc.models.agent_controls import AgentControls from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations +from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__name__) @@ -15,19 +17,29 @@ class Database(object): pass @staticmethod - def reset_db(): + def reset_db(reset_config=True): logger.info("Resetting database") # We can't drop system collections. [ Database.drop_collection(x) for x in mongo.db.collection_names() - if not x.startswith("system.") and not x == AttackMitigations.COLLECTION_NAME + if Database._should_drop(x, reset_config) ] ConfigService.init_config() Database.init_agent_controls() logger.info("DB was reset") return jsonify(status="OK") + @staticmethod + def _should_drop(collection: str, drop_config: bool) -> bool: + if not drop_config: + if collection == IslandMode.COLLECTION_NAME or collection == Config.COLLECTION_NAME: + return False + return ( + not collection.startswith("system.") + and not collection == AttackMitigations.COLLECTION_NAME + ) + @staticmethod def drop_collection(collection_name: str): mongo.db[collection_name].drop()