From bbcdc1bef4c41a3190b9c2a580ea89edcfdeeb19 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 5 Oct 2022 14:33:00 +0300 Subject: [PATCH 01/10] 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 From 6c913895c56014dc6e815e269e722321ab48b7fb Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 5 Oct 2022 14:46:47 +0300 Subject: [PATCH 02/10] Island: Add TCP connections to nodes based on TCP scan event --- .../scan_event_handler.py | 33 +++++++- .../test_scan_event_handler.py | 77 ++++++++++++++++--- 2 files changed, 96 insertions(+), 14 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 0eac3637c..1800eb737 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 @@ -1,12 +1,13 @@ +from copy import deepcopy from ipaddress import IPv4Interface from logging import getLogger from typing import Union from typing_extensions import TypeAlias -from common.agent_events import PingScanEvent, TCPScanEvent -from common.types import PortStatus -from monkey_island.cc.models import CommunicationType, Machine +from common.agent_events import AbstractAgentEvent, PingScanEvent, TCPScanEvent +from common.types import PortStatus, SocketAddress +from monkey_island.cc.models import CommunicationType, Machine, Node from monkey_island.cc.repository import ( IAgentRepository, IMachineRepository, @@ -56,11 +57,22 @@ class ScanEventHandler: try: target_machine = self._get_target_machine(event) + source_node = self._get_source_node(event) self._update_nodes(target_machine, event) + self._update_tcp_connections(source_node, target_machine, event) except (RetrievalError, StorageError, UnknownRecordError): logger.exception("Unable to process tcp scan data") + def _get_source_node(self, event: AbstractAgentEvent) -> Node: + machine = self._get_source_machine(event) + try: + return [ + node for node in self._node_repository.get_nodes() if node.machine_id == machine.id + ][0] + except KeyError: + raise UnknownRecordError(f"Source node for event {event} does not exist") + def _get_target_machine(self, event: ScanEvent) -> Machine: try: target_machines = self._machine_repository.get_machines_by_ip(event.target) @@ -85,6 +97,21 @@ class ScanEventHandler: src_machine.id, target_machine.id, CommunicationType.SCANNED ) + def _update_tcp_connections(self, src_node: Node, target_machine: Machine, event: TCPScanEvent): + node_connections = dict(deepcopy(src_node.tcp_connections)) + try: + machine_connections = set(node_connections[target_machine.id]) + except KeyError: + machine_connections = set() + open_ports = [port for port, status in event.ports.items() if status == PortStatus.OPEN] + for open_port in open_ports: + socket_address = SocketAddress(ip=event.target, port=open_port) + machine_connections.add(socket_address) + + node_connections[target_machine.id] = tuple(machine_connections) + src_node.tcp_connections = node_connections + self._node_repository.upsert_node(src_node) + def _get_source_machine(self, event: ScanEvent) -> Machine: agent = self._agent_repository.get_agent_by_id(event.source) return self._machine_repository.get_machine_by_id(agent.machine_id) 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 ad1ced7fa..f6832b788 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 @@ -1,3 +1,4 @@ +from copy import deepcopy from ipaddress import IPv4Address, IPv4Interface from itertools import count from unittest.mock import MagicMock @@ -9,7 +10,7 @@ from common import OperatingSystem from common.agent_events import PingScanEvent, TCPScanEvent from common.types import PortStatus, SocketAddress from monkey_island.cc.agent_event_handlers import ScanEventHandler -from monkey_island.cc.models import Agent, CommunicationType, Machine +from monkey_island.cc.models import Agent, CommunicationType, Machine, Node from monkey_island.cc.repository import ( IAgentRepository, IMachineRepository, @@ -29,43 +30,74 @@ SOURCE_MACHINE = Machine( hardware_id=5, network_interfaces=[IPv4Interface("10.10.10.99/24")], ) + +TARGET_MACHINE_ID = 33 +TARGET_MACHINE_IP = "10.10.10.1" TARGET_MACHINE = Machine( - id=33, + id=TARGET_MACHINE_ID, hardware_id=9, - network_interfaces=[IPv4Interface("10.10.10.1/24")], + network_interfaces=[IPv4Interface(f"{TARGET_MACHINE_IP}/24")], +) + +SOURCE_NODE = Node( + machine_id=SOURCE_MACHINE.id, + connections=[], + tcp_connections={ + 44: (SocketAddress(ip="1.1.1.1", port=40), SocketAddress(ip="2.2.2.2", port=50)) + }, +) + +SOURCE_NODE_2 = Node( + machine_id=SOURCE_MACHINE.id, + connections=[], + tcp_connections={ + 44: (SocketAddress(ip="1.1.1.1", port=40), SocketAddress(ip="2.2.2.2", port=50)), + TARGET_MACHINE_ID: (SocketAddress(ip=TARGET_MACHINE_IP, port=22),), + }, +) + +EXPECTED_NODE = Node( + machine_id=SOURCE_MACHINE.id, + connections=[], + tcp_connections={ + 44: (SocketAddress(ip="1.1.1.1", port=40), SocketAddress(ip="2.2.2.2", port=50)), + TARGET_MACHINE_ID: ( + SocketAddress(ip=TARGET_MACHINE_IP, port=22), + SocketAddress(ip=TARGET_MACHINE_IP, port=80), + ), + }, ) PING_SCAN_EVENT = PingScanEvent( source=AGENT_ID, - target=IPv4Address("10.10.10.1"), + target=IPv4Address(TARGET_MACHINE_IP), response_received=True, os=OperatingSystem.LINUX, ) PING_SCAN_EVENT_NO_RESPONSE = PingScanEvent( source=AGENT_ID, - target=IPv4Address("10.10.10.1"), + target=IPv4Address(TARGET_MACHINE_IP), response_received=False, os=OperatingSystem.LINUX, ) PING_SCAN_EVENT_NO_OS = PingScanEvent( source=AGENT_ID, - target=IPv4Address("10.10.10.1"), + target=IPv4Address(TARGET_MACHINE_IP), response_received=True, os=None, ) - TCP_SCAN_EVENT = TCPScanEvent( source=AGENT_ID, - target=IPv4Address("10.10.10.1"), - ports={22: PortStatus.OPEN, 8080: PortStatus.CLOSED}, + target=IPv4Address(TARGET_MACHINE_IP), + ports={22: PortStatus.OPEN, 80: PortStatus.OPEN, 8080: PortStatus.CLOSED}, ) TCP_SCAN_EVENT_CLOSED = TCPScanEvent( source=AGENT_ID, - target=IPv4Address("10.10.10.1"), + target=IPv4Address(TARGET_MACHINE_IP), ports={145: PortStatus.CLOSED, 8080: PortStatus.CLOSED}, ) @@ -91,6 +123,8 @@ def machine_repository() -> IMachineRepository: @pytest.fixture def node_repository() -> INodeRepository: node_repository = MagicMock(spec=INodeRepository) + node_repository.get_nodes.return_value = [deepcopy(SOURCE_NODE)] + node_repository.upsert_node = MagicMock() node_repository.upsert_communication = MagicMock() return node_repository @@ -103,7 +137,7 @@ def scan_event_handler(agent_repository, machine_repository, node_repository): MACHINES_BY_ID = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} MACHINES_BY_IP = { IPv4Address("10.10.10.99"): [SOURCE_MACHINE], - IPv4Address("10.10.10.1"): [TARGET_MACHINE], + IPv4Address(TARGET_MACHINE_IP): [TARGET_MACHINE], } @@ -186,6 +220,27 @@ def test_tcp_scan_event_target_machine_not_exists( machine_repository.upsert_machine.assert_called_with(expected_machine) +def test_handle_tcp_scan_event__tcp_connections( + scan_event_handler, machine_repository, node_repository +): + event = TCP_SCAN_EVENT + scan_event_handler._update_nodes = MagicMock() + scan_event_handler.handle_tcp_scan_event(event) + + node_repository.upsert_node.assert_called_with(EXPECTED_NODE) + + +def test_handle_tcp_scan_event__tcp_connections_upsert( + scan_event_handler, machine_repository, node_repository +): + event = TCP_SCAN_EVENT + node_repository.get_nodes.return_value = [deepcopy(SOURCE_NODE_2)] + scan_event_handler._update_nodes = MagicMock() + scan_event_handler.handle_tcp_scan_event(event) + + node_repository.upsert_node.assert_called_with(EXPECTED_NODE) + + @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], From 249950d602b28bfc9874dbf7d5329dad2fdc789b Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 5 Oct 2022 17:07:19 +0300 Subject: [PATCH 03/10] Island: Improve tcp handler code and coverage --- .../scan_event_handler.py | 7 ++---- .../test_scan_event_handler.py | 22 +++++++++++++++++-- 2 files changed, 22 insertions(+), 7 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 1800eb737..6f3a7030f 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 @@ -70,7 +70,7 @@ class ScanEventHandler: return [ node for node in self._node_repository.get_nodes() if node.machine_id == machine.id ][0] - except KeyError: + except IndexError: raise UnknownRecordError(f"Source node for event {event} does not exist") def _get_target_machine(self, event: ScanEvent) -> Machine: @@ -99,10 +99,7 @@ class ScanEventHandler: def _update_tcp_connections(self, src_node: Node, target_machine: Machine, event: TCPScanEvent): node_connections = dict(deepcopy(src_node.tcp_connections)) - try: - machine_connections = set(node_connections[target_machine.id]) - except KeyError: - machine_connections = set() + machine_connections = set(node_connections.get(target_machine.id, set())) open_ports = [port for port, status in event.ports.items() if status == PortStatus.OPEN] for open_port in open_ports: socket_address = SocketAddress(ip=event.target, port=open_port) 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 f6832b788..160d00ae1 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 @@ -227,7 +227,10 @@ def test_handle_tcp_scan_event__tcp_connections( scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) - node_repository.upsert_node.assert_called_with(EXPECTED_NODE) + node_passed = node_repository.upsert_node.call_args[0][0] + assert set(node_passed.tcp_connections[TARGET_MACHINE_ID]) == set( + EXPECTED_NODE.tcp_connections[TARGET_MACHINE_ID] + ) def test_handle_tcp_scan_event__tcp_connections_upsert( @@ -238,7 +241,22 @@ def test_handle_tcp_scan_event__tcp_connections_upsert( scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) - node_repository.upsert_node.assert_called_with(EXPECTED_NODE) + node_passed = node_repository.upsert_node.call_args[0][0] + assert set(node_passed.tcp_connections[TARGET_MACHINE_ID]) == set( + EXPECTED_NODE.tcp_connections[TARGET_MACHINE_ID] + ) + + +def test_handle_tcp_scan_event__no_source( + caplog, scan_event_handler, machine_repository, node_repository +): + event = TCP_SCAN_EVENT + node_repository.get_nodes.return_value = [] + scan_event_handler._update_nodes = MagicMock() + + scan_event_handler.handle_tcp_scan_event(event) + assert "ERROR" in caplog.text + assert f"Source node for event {event} does not exist" in caplog.text @pytest.mark.parametrize( From 2248bdcd67bb38561102b613f8c4718a3c57a241 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:10:47 +0300 Subject: [PATCH 04/10] Island: Add _get_node_by_id method to mongo_node_repository.py --- .../cc/repository/mongo_node_repository.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 447bfe8bb..8543d5e21 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -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}) From c90044074d78d81642682741cd9d7defdaacb274 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:21:13 +0300 Subject: [PATCH 05/10] Island: Remove storage error when node wasn't modified Upserting should throw an error when updating or inserting went wrong, not when a node is already up to date. --- .../cc/repository/mongo_node_repository.py | 6 ------ .../cc/repository/test_mongo_node_repository.py | 15 --------------- 2 files changed, 21 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 8543d5e21..248b0f973 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -55,12 +55,6 @@ class MongoNodeRepository(INodeRepository): except Exception as err: raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}") - if result.matched_count != 0 and result.modified_count != 1: - raise StorageError( - f'Error updating node with source ID "{node.machine_id}": Expected to update 1 ' - f"node, but {result.modified_count} were updated" - ) - if result.matched_count == 0 and result.upserted_id is None: raise StorageError( f'Error inserting node with source ID "{node.machine_id}": Expected to insert 1 ' 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 9d1ee6279..338526d76 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 @@ -163,21 +163,6 @@ def test_upsert_communication__replace_one_fails( error_raising_node_repository.upsert_communication(1, 2, CommunicationType.SCANNED) -def test_upsert_communication__replace_one_matched_without_modify( - error_raising_mock_mongo_client, error_raising_node_repository -): - mock_result = MagicMock() - mock_result.matched_count = 1 - mock_result.modified_count = 0 - error_raising_mock_mongo_client.monkey_island.nodes.find_one = MagicMock(return_value=None) - error_raising_mock_mongo_client.monkey_island.nodes.replace_one = MagicMock( - return_value=mock_result - ) - - with pytest.raises(StorageError): - error_raising_node_repository.upsert_communication(1, 2, CommunicationType.SCANNED) - - def test_upsert_communication__replace_one_insert_fails( error_raising_mock_mongo_client, error_raising_node_repository ): From b0ec035909501a1a5174e5b93ccb31989604eb35 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:31:12 +0300 Subject: [PATCH 06/10] Island: Move tcp_connection addition to node repository --- .../scan_event_handler.py | 15 +++---- monkey/monkey_island/cc/models/node.py | 5 ++- .../cc/repository/i_node_repository.py | 10 +++-- .../cc/repository/mongo_node_repository.py | 13 +++++- .../test_scan_event_handler.py | 43 +++++++++++-------- .../repository/test_mongo_node_repository.py | 39 +++++++++++++++-- 6 files changed, 87 insertions(+), 38 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 6f3a7030f..73e3023db 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 @@ -1,4 +1,3 @@ -from copy import deepcopy from ipaddress import IPv4Interface from logging import getLogger from typing import Union @@ -98,16 +97,16 @@ class ScanEventHandler: ) def _update_tcp_connections(self, src_node: Node, target_machine: Machine, event: TCPScanEvent): - node_connections = dict(deepcopy(src_node.tcp_connections)) - machine_connections = set(node_connections.get(target_machine.id, set())) - open_ports = [port for port, status in event.ports.items() if status == PortStatus.OPEN] + tcp_connections = set() + open_ports = (port for port, status in event.ports.items() if status == PortStatus.OPEN) for open_port in open_ports: socket_address = SocketAddress(ip=event.target, port=open_port) - machine_connections.add(socket_address) + tcp_connections.add(socket_address) - node_connections[target_machine.id] = tuple(machine_connections) - src_node.tcp_connections = node_connections - self._node_repository.upsert_node(src_node) + if tcp_connections: + self._node_repository.add_tcp_connections( + src_node.machine_id, {target_machine.id: tcp_connections} + ) def _get_source_machine(self, event: ScanEvent) -> Machine: agent = self._agent_repository.get_agent_by_id(event.source) diff --git a/monkey/monkey_island/cc/models/node.py b/monkey/monkey_island/cc/models/node.py index ada8aac19..95404f7d1 100644 --- a/monkey/monkey_island/cc/models/node.py +++ b/monkey/monkey_island/cc/models/node.py @@ -1,4 +1,4 @@ -from typing import FrozenSet, Mapping, Tuple +from typing import Dict, FrozenSet, Mapping, Tuple from pydantic import Field from typing_extensions import TypeAlias @@ -9,6 +9,7 @@ from common.types import SocketAddress from . import CommunicationType, MachineID NodeConnections: TypeAlias = Mapping[MachineID, FrozenSet[CommunicationType]] +TCPConnections: TypeAlias = Dict[MachineID, Tuple[SocketAddress, ...]] class Node(MutableInfectionMonkeyBaseModel): @@ -26,5 +27,5 @@ class Node(MutableInfectionMonkeyBaseModel): connections: NodeConnections """All outbound connections from this node to other machines""" - tcp_connections: Mapping[MachineID, Tuple[SocketAddress, ...]] = {} + tcp_connections: TCPConnections = {} """All successfull outbound 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 3738a1eb6..11983206c 100644 --- a/monkey/monkey_island/cc/repository/i_node_repository.py +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from typing import Sequence from monkey_island.cc.models import CommunicationType, MachineID, Node +from monkey_island.cc.models.node import TCPConnections class INodeRepository(ABC): @@ -26,11 +27,12 @@ class INodeRepository(ABC): """ @abstractmethod - def upsert_node(self, node: Node): + def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): """ - 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 + Add TCP connections to Node + :param machine_id: Machine ID of the Node that made the connections + :param tcp_connections: TCP connections made by node + :raises StorageError: If an error occurs while attempting to add connections """ @abstractmethod diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 248b0f973..599bddf61 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -5,6 +5,7 @@ from pymongo import MongoClient from monkey_island.cc.models import CommunicationType, MachineID, Node +from ..models.node import TCPConnections from . import INodeRepository, RemovalError, RetrievalError, StorageError from .consts import MONGO_OBJECT_ID_KEY @@ -47,7 +48,17 @@ class MongoNodeRepository(INodeRepository): return new_node - def upsert_node(self, node: Node): + def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): + node = self._get_node_by_id(machine_id) + + for target, connections in tcp_connections.items(): + if target in node.tcp_connections: + node.tcp_connections[target] = tuple({*node.tcp_connections[target], *connections}) + else: + node.tcp_connections[target] = connections + self._upsert_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 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 160d00ae1..1d8e71869 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 @@ -95,6 +95,13 @@ TCP_SCAN_EVENT = TCPScanEvent( ports={22: PortStatus.OPEN, 80: PortStatus.OPEN, 8080: PortStatus.CLOSED}, ) +TCP_CONNECTIONS = { + TARGET_MACHINE_ID: ( + SocketAddress(ip=TARGET_MACHINE_IP, port=22), + SocketAddress(ip=TARGET_MACHINE_IP, port=80), + ) +} + TCP_SCAN_EVENT_CLOSED = TCPScanEvent( source=AGENT_ID, target=IPv4Address(TARGET_MACHINE_IP), @@ -220,31 +227,29 @@ def test_tcp_scan_event_target_machine_not_exists( machine_repository.upsert_machine.assert_called_with(expected_machine) -def test_handle_tcp_scan_event__tcp_connections( +def test_handle_tcp_scan_event__no_open_ports( + scan_event_handler, machine_repository, node_repository +): + event = TCP_SCAN_EVENT_CLOSED + scan_event_handler._update_nodes = MagicMock() + scan_event_handler.handle_tcp_scan_event(event) + + assert not node_repository.add_tcp_connections.called + + +def test_handle_tcp_scan_event__ports_found( scan_event_handler, machine_repository, node_repository ): event = TCP_SCAN_EVENT scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) - node_passed = node_repository.upsert_node.call_args[0][0] - assert set(node_passed.tcp_connections[TARGET_MACHINE_ID]) == set( - EXPECTED_NODE.tcp_connections[TARGET_MACHINE_ID] - ) - - -def test_handle_tcp_scan_event__tcp_connections_upsert( - scan_event_handler, machine_repository, node_repository -): - event = TCP_SCAN_EVENT - node_repository.get_nodes.return_value = [deepcopy(SOURCE_NODE_2)] - scan_event_handler._update_nodes = MagicMock() - scan_event_handler.handle_tcp_scan_event(event) - - node_passed = node_repository.upsert_node.call_args[0][0] - assert set(node_passed.tcp_connections[TARGET_MACHINE_ID]) == set( - EXPECTED_NODE.tcp_connections[TARGET_MACHINE_ID] - ) + call_args = node_repository.add_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] + assert set(open_socket_addresses) == set(TCP_CONNECTIONS[TARGET_MACHINE_ID]) + assert len(open_socket_addresses) == len(TCP_CONNECTIONS[TARGET_MACHINE_ID]) def test_handle_tcp_scan_event__no_source( 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 338526d76..d4b16cd46 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 @@ -3,6 +3,7 @@ from unittest.mock import MagicMock import mongomock import pytest +from common.types import SocketAddress from monkey_island.cc.models import CommunicationType, Node from monkey_island.cc.repository import ( INodeRepository, @@ -12,6 +13,14 @@ from monkey_island.cc.repository import ( StorageError, ) +TARGET_MACHINE_IP = "2.2.2.2" + +TCP_CONNECTION_PORT_22 = {3: (SocketAddress(ip=TARGET_MACHINE_IP, port=22),)} +TCP_CONNECTION_PORT_80 = {3: (SocketAddress(ip=TARGET_MACHINE_IP, port=80),)} +ALL_TCP_CONNECTIONS = { + 3: (SocketAddress(ip=TARGET_MACHINE_IP, port=22), SocketAddress(ip=TARGET_MACHINE_IP, port=80)) +} + NODES = ( Node( machine_id=1, @@ -23,6 +32,7 @@ NODES = ( Node( machine_id=2, connections={1: frozenset((CommunicationType.CC,))}, + tcp_connections=TCP_CONNECTION_PORT_22, ), Node( machine_id=3, @@ -32,10 +42,7 @@ NODES = ( 5: frozenset((CommunicationType.SCANNED, CommunicationType.EXPLOITED)), }, ), - Node( - machine_id=4, - connections={}, - ), + Node(machine_id=4, connections={}, tcp_connections=ALL_TCP_CONNECTIONS), Node( machine_id=5, connections={ @@ -201,3 +208,27 @@ def test_reset(node_repository): def test_reset__removal_error(error_raising_node_repository): with pytest.raises(RemovalError): error_raising_node_repository.reset() + + +def test_upsert_tcp_connections__empty_connections(node_repository): + node_repository.add_tcp_connections(1, TCP_CONNECTION_PORT_22) + nodes = node_repository.get_nodes() + for node in nodes: + if node.machine_id == 1: + assert node.tcp_connections == TCP_CONNECTION_PORT_22 + + +def test_upsert_tcp_connections__upsert_new_port(node_repository): + node_repository.add_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) + assert len(modified_node.tcp_connections) == len(ALL_TCP_CONNECTIONS) + + +def test_upsert_tcp_connections__port_already_present(node_repository): + node_repository.add_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) + assert len(modified_node.tcp_connections) == len(ALL_TCP_CONNECTIONS) From 3bc2e4876fac15f34ea86b7237e8248863e95fff Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:45:56 +0300 Subject: [PATCH 07/10] Island: Handle missing node in add_tcp_connections --- .../monkey_island/cc/repository/mongo_node_repository.py | 3 +++ .../cc/repository/test_mongo_node_repository.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index 599bddf61..b417a7554 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -51,6 +51,9 @@ class MongoNodeRepository(INodeRepository): def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): node = self._get_node_by_id(machine_id) + if node is None: + node = Node(machine_id=machine_id, connections={}) + for target, connections in tcp_connections.items(): if target in node.tcp_connections: node.tcp_connections[target] = tuple({*node.tcp_connections[target], *connections}) 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 d4b16cd46..3b2a0d26f 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 @@ -232,3 +232,10 @@ def test_upsert_tcp_connections__port_already_present(node_repository): modified_node = [node for node in nodes if node.machine_id == 4][0] assert set(modified_node.tcp_connections) == set(ALL_TCP_CONNECTIONS) assert len(modified_node.tcp_connections) == len(ALL_TCP_CONNECTIONS) + + +def test_upsert_tcp_connections__node_missing(node_repository): + node_repository.add_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) From 0d246a04793212b1f1d256e95b75536e5c603e21 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 14:49:11 +0300 Subject: [PATCH 08/10] 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) From 8503e0f49964f27b8ffcc6fa83606b12d6681b53 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Thu, 6 Oct 2022 15:02:13 +0300 Subject: [PATCH 09/10] UT: Remove unused test data structures --- .../test_scan_event_handler.py | 21 ------------------- 1 file changed, 21 deletions(-) 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 280d73e55..e00896e3f 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 @@ -47,27 +47,6 @@ SOURCE_NODE = Node( }, ) -SOURCE_NODE_2 = Node( - machine_id=SOURCE_MACHINE.id, - connections=[], - tcp_connections={ - 44: (SocketAddress(ip="1.1.1.1", port=40), SocketAddress(ip="2.2.2.2", port=50)), - TARGET_MACHINE_ID: (SocketAddress(ip=TARGET_MACHINE_IP, port=22),), - }, -) - -EXPECTED_NODE = Node( - machine_id=SOURCE_MACHINE.id, - connections=[], - tcp_connections={ - 44: (SocketAddress(ip="1.1.1.1", port=40), SocketAddress(ip="2.2.2.2", port=50)), - TARGET_MACHINE_ID: ( - SocketAddress(ip=TARGET_MACHINE_IP, port=22), - SocketAddress(ip=TARGET_MACHINE_IP, port=80), - ), - }, -) - PING_SCAN_EVENT = PingScanEvent( source=AGENT_ID, target=IPv4Address(TARGET_MACHINE_IP), From be4ecccdcdef0104d2aa5ad6a459b4bf604df53f Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 7 Oct 2022 10:05:06 +0300 Subject: [PATCH 10/10] Island: Refactor get_node_by_id to raise UnknownRecordError --- .../scan_event_handler.py | 7 +---- .../cc/repository/i_node_repository.py | 9 ++++++ .../cc/repository/mongo_node_repository.py | 30 +++++++++---------- .../test_scan_event_handler.py | 5 ++-- .../repository/test_mongo_node_repository.py | 10 +++++++ 5 files changed, 38 insertions(+), 23 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 8c769a470..f3482b891 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 @@ -65,12 +65,7 @@ class ScanEventHandler: def _get_source_node(self, event: AbstractAgentEvent) -> Node: machine = self._get_source_machine(event) - try: - return [ - node for node in self._node_repository.get_nodes() if node.machine_id == machine.id - ][0] - except IndexError: - raise UnknownRecordError(f"Source node for event {event} does not exist") + return self._node_repository.get_node_by_machine_id(machine.id) def _get_target_machine(self, event: ScanEvent) -> Machine: try: diff --git a/monkey/monkey_island/cc/repository/i_node_repository.py b/monkey/monkey_island/cc/repository/i_node_repository.py index 181cd185e..157c9274c 100644 --- a/monkey/monkey_island/cc/repository/i_node_repository.py +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -44,6 +44,15 @@ class INodeRepository(ABC): :raises RetrievalError: If an error occurs while attempting to retrieve the nodes """ + @abstractmethod + def get_node_by_machine_id(self, machine_id: MachineID) -> Node: + """ + Fetches network Node from the database based on Machine id + :param machine_id: ID of a Machine that Node represents + :return: network Node that represents the Machine + :raises UnknownRecordError: If the Node does not exist + """ + @abstractmethod def reset(self): """ diff --git a/monkey/monkey_island/cc/repository/mongo_node_repository.py b/monkey/monkey_island/cc/repository/mongo_node_repository.py index b406b8fad..befc81632 100644 --- a/monkey/monkey_island/cc/repository/mongo_node_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_node_repository.py @@ -1,12 +1,12 @@ from copy import deepcopy -from typing import Optional, Sequence +from typing import Sequence from pymongo import MongoClient from monkey_island.cc.models import CommunicationType, MachineID, Node from ..models.node import TCPConnections -from . import INodeRepository, RemovalError, RetrievalError, StorageError +from . import INodeRepository, RemovalError, RetrievalError, StorageError, UnknownRecordError from .consts import MONGO_OBJECT_ID_KEY UPSERT_ERROR_MESSAGE = "An error occurred while attempting to upsert a node" @@ -21,16 +21,14 @@ class MongoNodeRepository(INodeRepository): self, src: MachineID, dst: MachineID, communication_type: CommunicationType ): try: - node = self._get_node_by_id(src) - except Exception as err: - raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}") - - if node is None: - updated_node = Node(machine_id=src, connections={dst: frozenset((communication_type,))}) - else: + node = self.get_node_by_machine_id(src) updated_node = MongoNodeRepository._add_connection_to_node( node, dst, communication_type ) + except UnknownRecordError: + updated_node = Node(machine_id=src, connections={dst: frozenset((communication_type,))}) + except Exception as err: + raise StorageError(f"{UPSERT_ERROR_MESSAGE}: {err}") self._upsert_node(updated_node) @@ -49,9 +47,9 @@ class MongoNodeRepository(INodeRepository): return new_node def upsert_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections): - node = self._get_node_by_id(machine_id) - - if node is None: + try: + node = self.get_node_by_machine_id(machine_id) + except UnknownRecordError: node = Node(machine_id=machine_id, connections={}) for target, connections in tcp_connections.items(): @@ -75,11 +73,13 @@ class MongoNodeRepository(INodeRepository): f"node, but no nodes were inserted" ) - def _get_node_by_id(self, node_id: MachineID) -> Optional[Node]: + def get_node_by_machine_id(self, machine_id: MachineID) -> Node: node_dict = self._nodes_collection.find_one( - {SRC_FIELD_NAME: node_id}, {MONGO_OBJECT_ID_KEY: False} + {SRC_FIELD_NAME: machine_id}, {MONGO_OBJECT_ID_KEY: False} ) - return Node(**node_dict) if node_dict else None + if not node_dict: + raise UnknownRecordError(f"Node with machine ID {machine_id}") + return Node(**node_dict) def get_nodes(self) -> Sequence[Node]: try: 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 e00896e3f..673f8293c 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 @@ -221,6 +221,7 @@ def test_handle_tcp_scan_event__ports_found( ): event = TCP_SCAN_EVENT scan_event_handler._update_nodes = MagicMock() + node_repository.get_node_by_machine_id.return_value = SOURCE_NODE scan_event_handler.handle_tcp_scan_event(event) call_args = node_repository.upsert_tcp_connections.call_args[0] @@ -235,12 +236,12 @@ def test_handle_tcp_scan_event__no_source( caplog, scan_event_handler, machine_repository, node_repository ): event = TCP_SCAN_EVENT - node_repository.get_nodes.return_value = [] + node_repository.get_node_by_machine_id = MagicMock(side_effect=UnknownRecordError("no source")) scan_event_handler._update_nodes = MagicMock() scan_event_handler.handle_tcp_scan_event(event) assert "ERROR" in caplog.text - assert f"Source node for event {event} does not exist" in caplog.text + assert "no source" in caplog.text @pytest.mark.parametrize( 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 fa95043d6..bf4968406 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 @@ -11,6 +11,7 @@ from monkey_island.cc.repository import ( RemovalError, RetrievalError, StorageError, + UnknownRecordError, ) TARGET_MACHINE_IP = "2.2.2.2" @@ -239,3 +240,12 @@ def test_upsert_tcp_connections__node_missing(node_repository): 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) + + +def test_get_node_by_machine_id(node_repository): + assert node_repository.get_node_by_machine_id(1) == NODES[0] + + +def test_get_node_by_machine_id__no_node(node_repository): + with pytest.raises(UnknownRecordError): + node_repository.get_node_by_machine_id(999)