From 0d246a04793212b1f1d256e95b75536e5c603e21 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:49:11 +0300 Subject: [PATCH] Island: Rename add_tcp_connections to upsert_tcp_connections --- .../cc/agent_event_handlers/scan_event_handler.py | 2 +- monkey/monkey_island/cc/repository/i_node_repository.py | 2 +- .../monkey_island/cc/repository/mongo_node_repository.py | 2 +- .../cc/agent_event_handlers/test_scan_event_handler.py | 4 ++-- .../cc/repository/test_mongo_node_repository.py | 8 ++++---- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py b/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py index 73e3023db..8c769a470 100644 --- a/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py +++ b/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py @@ -104,7 +104,7 @@ class ScanEventHandler: tcp_connections.add(socket_address) if tcp_connections: - self._node_repository.add_tcp_connections( + self._node_repository.upsert_tcp_connections( src_node.machine_id, {target_machine.id: tcp_connections} ) diff --git a/monkey/monkey_island/cc/repository/i_node_repository.py b/monkey/monkey_island/cc/repository/i_node_repository.py index 11983206c..181cd185e 100644 --- a/monkey/monkey_island/cc/repository/i_node_repository.py +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -27,7 +27,7 @@ class INodeRepository(ABC): """ @abstractmethod - def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): + def upsert_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): """ Add TCP connections to Node :param machine_id: Machine ID of the Node that made the connections diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index b417a7554..b406b8fad 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -48,7 +48,7 @@ class MongoNodeRepository(INodeRepository): return new_node - def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): + def upsert_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): node = self._get_node_by_id(machine_id) if node is None: diff --git a/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py index 1d8e71869..280d73e55 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py +++ b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py @@ -234,7 +234,7 @@ def test_handle_tcp_scan_event__no_open_ports( scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) - assert not node_repository.add_tcp_connections.called + assert not node_repository.upsert_tcp_connections.called def test_handle_tcp_scan_event__ports_found( @@ -244,7 +244,7 @@ def test_handle_tcp_scan_event__ports_found( scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) - call_args = node_repository.add_tcp_connections.call_args[0] + call_args = node_repository.upsert_tcp_connections.call_args[0] assert call_args[0] == MACHINE_ID assert TARGET_MACHINE_ID in call_args[1] open_socket_addresses = call_args[1][TARGET_MACHINE_ID] diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_node_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_node_repository.py index 3b2a0d26f..fa95043d6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_node_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_node_repository.py @@ -211,7 +211,7 @@ def test_reset__removal_error(error_raising_node_repository): def test_upsert_tcp_connections__empty_connections(node_repository): - node_repository.add_tcp_connections(1, TCP_CONNECTION_PORT_22) + node_repository.upsert_tcp_connections(1, TCP_CONNECTION_PORT_22) nodes = node_repository.get_nodes() for node in nodes: if node.machine_id == 1: @@ -219,7 +219,7 @@ def test_upsert_tcp_connections__empty_connections(node_repository): def test_upsert_tcp_connections__upsert_new_port(node_repository): - node_repository.add_tcp_connections(2, TCP_CONNECTION_PORT_80) + node_repository.upsert_tcp_connections(2, TCP_CONNECTION_PORT_80) nodes = node_repository.get_nodes() modified_node = [node for node in nodes if node.machine_id == 2][0] assert set(modified_node.tcp_connections) == set(ALL_TCP_CONNECTIONS) @@ -227,7 +227,7 @@ def test_upsert_tcp_connections__upsert_new_port(node_repository): def test_upsert_tcp_connections__port_already_present(node_repository): - node_repository.add_tcp_connections(4, TCP_CONNECTION_PORT_80) + node_repository.upsert_tcp_connections(4, TCP_CONNECTION_PORT_80) nodes = node_repository.get_nodes() modified_node = [node for node in nodes if node.machine_id == 4][0] assert set(modified_node.tcp_connections) == set(ALL_TCP_CONNECTIONS) @@ -235,7 +235,7 @@ def test_upsert_tcp_connections__port_already_present(node_repository): def test_upsert_tcp_connections__node_missing(node_repository): - node_repository.add_tcp_connections(999, TCP_CONNECTION_PORT_80) + node_repository.upsert_tcp_connections(999, TCP_CONNECTION_PORT_80) nodes = node_repository.get_nodes() modified_node = [node for node in nodes if node.machine_id == 999][0] assert set(modified_node.tcp_connections) == set(TCP_CONNECTION_PORT_80)