From bbcdc1bef4c41a3190b9c2a580ea89edcfdeeb19 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 5 Oct 2022 14:33:00 +0300 Subject: [PATCH] Island: Make upsert_node method public Updating/inserting the node into the repository is required outside of repository itself. --- monkey/monkey_island/cc/repository/i_node_repository.py | 8 ++++++++ .../monkey_island/cc/repository/mongo_node_repository.py | 4 ++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_node_repository.py b/monkey/monkey_island/cc/repository/i_node_repository.py index 53fc53bc9..3738a1eb6 100644 --- a/monkey/monkey_island/cc/repository/i_node_repository.py +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -25,6 +25,14 @@ class INodeRepository(ABC): :raises StorageError: If an error occurs while attempting to upsert the Node """ + @abstractmethod + def upsert_node(self, node: Node): + """ + Store the Node object in the repository by creating a new one or updating an existing one. + :param node: Node that will be saved + :raises StorageError: If an error occurs while attempting to upsert the Node + """ + @abstractmethod def get_nodes(self) -> Sequence[Node]: """ diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 9b994d110..447bfe8bb 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -34,7 +34,7 @@ class MongoNodeRepository(INodeRepository): node, dst, communication_type ) - self._upsert_node(updated_node) + self.upsert_node(updated_node) @staticmethod def _add_connection_to_node( @@ -50,7 +50,7 @@ class MongoNodeRepository(INodeRepository): return new_node - def _upsert_node(self, node: Node): + def upsert_node(self, node: Node): try: result = self._nodes_collection.replace_one( {SRC_FIELD_NAME: node.machine_id}, node.dict(simplify=True), upsert=True