forked from p15670423/monkey
Island: Omit the mongo object ID from Machine query results
This commit is contained in:
parent
3fb1ddaa74
commit
fa8736e2fe
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue