p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 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 typing import Sequence
from typing import Optional, Sequence
from pymongo import MongoClient
@ -20,21 +20,18 @@ class MongoNodeRepository(INodeRepository):
self, src: MachineID, dst: MachineID, communication_type: CommunicationType
):
try:
node_dict = self._nodes_collection.find_one(
{SRC_FIELD_NAME: src}, {MONGO_OBJECT_ID_KEY: False}
)
node = self._get_node_by_id(src)
except Exception as 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,))})
else:
node = Node(**node_dict)
updated_node = MongoNodeRepository._add_connection_to_node(
node, dst, communication_type
)
self.upsert_node(updated_node)
self._upsert_node(updated_node)
@staticmethod
def _add_connection_to_node(
@ -70,6 +67,12 @@ class MongoNodeRepository(INodeRepository):
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]:
try:
cursor = self._nodes_collection.find({}, {MONGO_OBJECT_ID_KEY: False})