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/vulture_allowlist.py b/vulture_allowlist.py index c77390784..c4c33b662 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 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,7 @@ ICredentialsRepository.save_configured_credentials ITelemetryRepository.get_telemetries IEventRepository.get_events IFindingRepository.get_findings +MongoMachineRepository key_list simulation netmap