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

Island: Omit the mongo object ID from Node query results

This commit is contained in:
Mike Salvatore 2022-09-14 09:02:21 -04:00
parent 5821d71ade
commit 3fb1ddaa74
1 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,5 +1,5 @@
from copy import deepcopy from copy import deepcopy
from typing import Any, MutableMapping, Sequence from typing import Sequence
from pymongo import MongoClient from pymongo import MongoClient
@ -20,25 +20,22 @@ class MongoNodeRepository(INodeRepository):
self, src: MachineID, dst: MachineID, communication_type: CommunicationType self, src: MachineID, dst: MachineID, communication_type: CommunicationType
): ):
try: try:
node_dict = self._nodes_collection.find_one({SRC_FIELD_NAME: src}) node_dict = self._nodes_collection.find_one(
{SRC_FIELD_NAME: src}, {MONGO_OBJECT_ID_KEY: False}
)
except Exception as err: except Exception as err:
raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}") raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}")
if node_dict is None: if node_dict is None:
updated_node = Node(machine_id=src, connections={dst: frozenset((communication_type,))}) updated_node = Node(machine_id=src, connections={dst: frozenset((communication_type,))})
else: else:
node = MongoNodeRepository._mongo_record_to_node(node_dict) node = Node(**node_dict)
updated_node = MongoNodeRepository._add_connection_to_node( updated_node = MongoNodeRepository._add_connection_to_node(
node, dst, communication_type node, dst, communication_type
) )
self._upsert_node(updated_node) self._upsert_node(updated_node)
@staticmethod
def _mongo_record_to_node(mongo_record: MutableMapping[str, Any]) -> Node:
del mongo_record[MONGO_OBJECT_ID_KEY]
return Node(**mongo_record)
@staticmethod @staticmethod
def _add_connection_to_node( def _add_connection_to_node(
node: Node, dst: MachineID, communication_type: CommunicationType node: Node, dst: MachineID, communication_type: CommunicationType
@ -75,8 +72,8 @@ class MongoNodeRepository(INodeRepository):
def get_nodes(self) -> Sequence[Node]: def get_nodes(self) -> Sequence[Node]:
try: try:
cursor = self._nodes_collection.find() cursor = self._nodes_collection.find({}, {MONGO_OBJECT_ID_KEY: False})
return list(map(MongoNodeRepository._mongo_record_to_node, cursor)) return list(map(lambda n: Node(**n), cursor))
except Exception as err: except Exception as err:
raise RetrievalError(f"Error retrieving nodes from the repository: {err}") raise RetrievalError(f"Error retrieving nodes from the repository: {err}")