Merge branch 'omit-mongo-object-id' into 2255-mongo-node-repository

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

monkey/monkey_island/cc/repository/mongo_credentials_repository.py

@@ -55,9 +55,8 @@ class MongoCredentialsRepository(ICredentialsRepository):
     def _get_credentials_from_collection(self, collection) -> Sequence[Credentials]:
         try:
             collection_result = []
-            list_collection_result = list(collection.find({}))
+            list_collection_result = list(collection.find({}, {MONGO_OBJECT_ID_KEY: False}))
             for encrypted_credentials in list_collection_result:
-                del encrypted_credentials[MONGO_OBJECT_ID_KEY]
                 plaintext_credentials = self._decrypt_credentials_mapping(encrypted_credentials)
                 collection_result.append(Credentials(**plaintext_credentials))
 

monkey/monkey_island/cc/repository/mongo_machine_repository.py

@@ -1,6 +1,6 @@
 from ipaddress import IPv4Address
 from threading import Lock
-from typing import Any, MutableMapping, Sequence
+from typing import Any, Sequence
 
 from pymongo import MongoClient
 
@@ -58,36 +58,33 @@ class MongoMachineRepository(IMachineRepository):
 
     def _find_one(self, key: str, search_value: Any) -> Machine:
         try:
-            machine_dict = self._machines_collection.find_one({key: search_value})
+            machine_dict = self._machines_collection.find_one(
+                {key: search_value}, {MONGO_OBJECT_ID_KEY: False}
+            )
         except Exception as err:
             raise RetrievalError(f'Error retrieving machine with "{key} == {search_value}": {err}')
 
         if machine_dict is None:
             raise UnknownRecordError(f'Unknown machine with "{key} == {search_value}"')
 
-        return MongoMachineRepository._mongo_record_to_machine(machine_dict)
+        return Machine(**machine_dict)
 
     def get_machines_by_ip(self, ip: IPv4Address) -> Sequence[Machine]:
         ip_regex = "^" + str(ip).replace(".", "\\.") + "\\/.*$"
         query = {"network_interfaces": {"$elemMatch": {"$regex": ip_regex}}}
 
         try:
-            cursor = self._machines_collection.find(query)
+            cursor = self._machines_collection.find(query, {MONGO_OBJECT_ID_KEY: False})
         except Exception as err:
             raise RetrievalError(f'Error retrieving machines with ip "{ip}": {err}')
 
-        machines = list(map(MongoMachineRepository._mongo_record_to_machine, cursor))
+        machines = list(map(lambda m: Machine(**m), cursor))
 
         if len(machines) == 0:
             raise UnknownRecordError(f'No machines found with IP "{ip}"')
 
         return machines
 
-    @staticmethod
-    def _mongo_record_to_machine(mongo_record: MutableMapping[str, Any]) -> Machine:
-        del mongo_record[MONGO_OBJECT_ID_KEY]
-        return Machine(**mongo_record)
-
     def reset(self):
         try:
             self._machines_collection.drop()

monkey/monkey_island/cc/repository/mongo_node_repository.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Any, MutableMapping, Sequence
+from typing import Sequence
 
 from pymongo import MongoClient
 
@@ -20,25 +20,22 @@ class MongoNodeRepository(INodeRepository):
         self, src: MachineID, dst: MachineID, communication_type: CommunicationType
     ):
         try:
-            node_dict = self._nodes_collection.find_one({SRC_FIELD_NAME: src})
+            node_dict = self._nodes_collection.find_one(
+                {SRC_FIELD_NAME: src}, {MONGO_OBJECT_ID_KEY: False}
+            )
         except Exception as err:
             raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}")
 
         if node_dict is None:
             updated_node = Node(machine_id=src, connections={dst: frozenset((communication_type,))})
         else:
-            node = MongoNodeRepository._mongo_record_to_node(node_dict)
+            node = Node(**node_dict)
             updated_node = MongoNodeRepository._add_connection_to_node(
                 node, dst, communication_type
             )
 
         self._upsert_node(updated_node)
 
-    @staticmethod
-    def _mongo_record_to_node(mongo_record: MutableMapping[str, Any]) -> Node:
-        del mongo_record[MONGO_OBJECT_ID_KEY]
-        return Node(**mongo_record)
-
     @staticmethod
     def _add_connection_to_node(
         node: Node, dst: MachineID, communication_type: CommunicationType
@@ -75,8 +72,8 @@ class MongoNodeRepository(INodeRepository):
 
     def get_nodes(self) -> Sequence[Node]:
         try:
-            cursor = self._nodes_collection.find()
-            return list(map(MongoNodeRepository._mongo_record_to_node, cursor))
+            cursor = self._nodes_collection.find({}, {MONGO_OBJECT_ID_KEY: False})
+            return list(map(lambda n: Node(**n), cursor))
         except Exception as err:
             raise RetrievalError(f"Error retrieving nodes from the repository: {err}")
 

vulture_allowlist.py

@@ -12,6 +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 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
@@ -277,6 +278,8 @@ ICredentialsRepository.save_configured_credentials
 ITelemetryRepository.get_telemetries
 IEventRepository.get_events
 IFindingRepository.get_findings
+MongoAgentRepository
+MongoMachineRepository
 key_list
 simulation
 netmap