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 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()
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue