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

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.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