Island: Omit the mongo object ID from Machine query results

This commit is contained in:
Mike Salvatore 2022-09-14 09:05:42 -04:00
parent 3fb1ddaa74
commit fa8736e2fe
2 changed files with 9 additions and 10 deletions

View File

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

View File

@ -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.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
from monkey_island.cc.models import Report from monkey_island.cc.models import Report
from monkey_island.cc.models.networkmap import Arc, NetworkMap 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.attack.IMitigationsRepository import IMitigationsRepository
from monkey_island.cc.repository.i_agent_repository import IAgentRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository
from monkey_island.cc.repository.i_attack_repository import IAttackRepository from monkey_island.cc.repository.i_attack_repository import IAttackRepository
@ -277,6 +278,7 @@ ICredentialsRepository.save_configured_credentials
ITelemetryRepository.get_telemetries ITelemetryRepository.get_telemetries
IEventRepository.get_events IEventRepository.get_events
IFindingRepository.get_findings IFindingRepository.get_findings
MongoMachineRepository
key_list key_list
simulation simulation
netmap netmap