Island: Omit the mongo object ID from Agent query results

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()

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