p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Add _get_node_by_id method to mongo_node_repository.py

This commit is contained in:
vakarisz 2022-10-06 14:10:47 +03:00
parent 249950d602
commit 2248bdcd67
1 changed files with 10 additions and 7 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 Sequence from typing import Optional, Sequence
from pymongo import MongoClient from pymongo import MongoClient
@ -20,21 +20,18 @@ 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( node = self._get_node_by_id(src)
{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 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 = 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 @staticmethod
def _add_connection_to_node( def _add_connection_to_node(
@ -70,6 +67,12 @@ class MongoNodeRepository(INodeRepository):
f"node, but no nodes were inserted" f"node, but no nodes were inserted"
) )
def _get_node_by_id(self, node_id: MachineID) -> Optional[Node]:
node_dict = self._nodes_collection.find_one(
{SRC_FIELD_NAME: node_id}, {MONGO_OBJECT_ID_KEY: False}
)
return Node(**node_dict) if node_dict else None
def get_nodes(self) -> Sequence[Node]: def get_nodes(self) -> Sequence[Node]:
try: try:
cursor = self._nodes_collection.find({}, {MONGO_OBJECT_ID_KEY: False}) cursor = self._nodes_collection.find({}, {MONGO_OBJECT_ID_KEY: False})