diff --git a/monkey/monkey_island/cc/repository/mongo_agent_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_repository.py index 3b1b60f0a..dfad2bbf7 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_repository.py @@ -1,4 +1,4 @@ -from typing import Any, MutableMapping, Sequence +from typing import Sequence from pymongo import MongoClient @@ -40,27 +40,24 @@ class MongoAgentRepository(IAgentRepository): def get_agent_by_id(self, agent_id: AgentID) -> Agent: try: - agent_dict = self._agents_collection.find_one({"id": str(agent_id)}) + agent_dict = self._agents_collection.find_one( + {"id": str(agent_id)}, {MONGO_OBJECT_ID_KEY: False} + ) except Exception as err: raise RetrievalError(f'Error retrieving agent with "id == {agent_id}": {err}') if agent_dict is None: raise UnknownRecordError(f'Unknown ID "{agent_id}"') - return MongoAgentRepository._mongo_record_to_agent(agent_dict) + return Agent(**agent_dict) def get_running_agents(self) -> Sequence[Agent]: try: - cursor = self._agents_collection.find({"stop_time": None}) - return list(map(MongoAgentRepository._mongo_record_to_agent, cursor)) + cursor = self._agents_collection.find({"stop_time": None}, {MONGO_OBJECT_ID_KEY: False}) + return list(map(lambda a: Agent(**a), cursor)) except Exception as err: raise RetrievalError(f"Error retrieving running agents: {err}") - @staticmethod - def _mongo_record_to_agent(mongo_record: MutableMapping[str, Any]) -> Agent: - del mongo_record[MONGO_OBJECT_ID_KEY] - return Agent(**mongo_record) - def reset(self): try: self._agents_collection.drop() diff --git a/vulture_allowlist.py b/vulture_allowlist.py index c4c33b662..98db69737 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -12,7 +12,7 @@ from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFacto from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap -from monkey_island.cc.repository import MongoMachineRepository +from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository from monkey_island.cc.repository.i_attack_repository import IAttackRepository @@ -278,6 +278,7 @@ ICredentialsRepository.save_configured_credentials ITelemetryRepository.get_telemetries IEventRepository.get_events IFindingRepository.get_findings +MongoAgentRepository MongoMachineRepository key_list simulation