p15693087
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests 1 Projects Releases Wiki Activity

Island: Omit the mongo object ID from Agent query results

This commit is contained in:
Mike Salvatore 2022-09-14 09:07:39 -04:00
parent fa8736e2fe
commit c2437464c6
2 changed files with 9 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
monkey/monkey_island/cc/repository/mongo_agent_repository.py
View File

@ -1,4 +1,4 @@
from typing import Any, MutableMapping, Sequence from typing import Sequence
from pymongo import MongoClient from pymongo import MongoClient
@ -40,27 +40,24 @@ class MongoAgentRepository(IAgentRepository):
def get_agent_by_id(self, agent_id: AgentID) -> Agent: def get_agent_by_id(self, agent_id: AgentID) -> Agent:
try: try:
agent_dict = self._agents_collection.find_one({"id": str(agent_id)}) agent_dict = self._agents_collection.find_one(
{"id": str(agent_id)}, {MONGO_OBJECT_ID_KEY: False}
)
except Exception as err: except Exception as err:
raise RetrievalError(f'Error retrieving agent with "id == {agent_id}": {err}') raise RetrievalError(f'Error retrieving agent with "id == {agent_id}": {err}')
if agent_dict is None: if agent_dict is None:
raise UnknownRecordError(f'Unknown ID "{agent_id}"') raise UnknownRecordError(f'Unknown ID "{agent_id}"')
return MongoAgentRepository._mongo_record_to_agent(agent_dict) return Agent(**agent_dict)
def get_running_agents(self) -> Sequence[Agent]: def get_running_agents(self) -> Sequence[Agent]:
try: try:
cursor = self._agents_collection.find({"stop_time": None}) cursor = self._agents_collection.find({"stop_time": None}, {MONGO_OBJECT_ID_KEY: False})
return list(map(MongoAgentRepository._mongo_record_to_agent, cursor)) return list(map(lambda a: Agent(**a), cursor))
except Exception as err: except Exception as err:
raise RetrievalError(f"Error retrieving running agents: {err}") raise RetrievalError(f"Error retrieving running agents: {err}")
@staticmethod
def _mongo_record_to_agent(mongo_record: MutableMapping[str, Any]) -> Agent:
del mongo_record[MONGO_OBJECT_ID_KEY]
return Agent(**mongo_record)
def reset(self): def reset(self):
try: try:
self._agents_collection.drop() self._agents_collection.drop()

3
vulture_allowlist.py
View File

@ -12,7 +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 import MongoAgentRepository, 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
@ -278,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
MongoAgentRepository
MongoMachineRepository MongoMachineRepository
key_list key_list
simulation simulation