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/monkey/monkey_island/cc/repository/mongo_credentials_repository.py b/monkey/monkey_island/cc/repository/mongo_credentials_repository.py index 237f052e4..3fdc306a8 100644 --- a/monkey/monkey_island/cc/repository/mongo_credentials_repository.py +++ b/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)) diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index 4d1a36470..fab038694 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/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() diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 78550d153..5b1b7f71f 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/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}") diff --git a/vulture_allowlist.py b/vulture_allowlist.py index c77390784..98db69737 100644 --- a/vulture_allowlist.py +++ b/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