From e4aec8b9a3dcd8038a16cbe7319d06b8d78d568a Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:15:50 +0200 Subject: [PATCH 01/24] Island: Add ScanEventHandler to agent_event_handlers --- .../scan_event_handler.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py 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 new file mode 100644 index 000000000..2d1da09e6 --- /dev/null +++ b/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py @@ -0,0 +1,91 @@ +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 monkey_island.cc.repository import ( + IAgentRepository, + IMachineRepository, + INodeRepository, + RetrievalError, + StorageError, + UnknownRecordError, +) + +ScanEvent: TypeAlias = Union[PingScanEvent, TCPScanEvent] + +logger = getLogger(__name__) + + +class ScanEventHandler: + """ + Handles scan event and makes changes to Machine and Node states base on it + """ + + def __init__( + self, + agent_repository: IAgentRepository, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + ): + self._agent_repository = agent_repository + self._machine_repository = machine_repository + self._node_repository = node_repository + + def handle_ping_scan_event(self, event: PingScanEvent): + if not event.response_received: + return + + try: + target_machine = self._get_target_machine(event) + + self._update_target_machine_os(target_machine, event) + self._update_nodes(target_machine, event) + except (RetrievalError, StorageError, TypeError, UnknownRecordError): + logger.exception("Unable to process ping scan data") + + def handle_tcp_scan_event(self, event: TCPScanEvent): + open_ports = [port for port, status in event.ports.items() if status == PortStatus.OPEN] + + if not open_ports: + return + + try: + target_machine = self._get_target_machine(event) + + self._update_nodes(target_machine, event) + except (RetrievalError, StorageError, UnknownRecordError): + logger.exception("Unable to process tcp scan data") + + def _get_target_machine(self, event: ScanEvent) -> Machine: + try: + target_machines = self._machine_repository.get_machines_by_ip(event.target) + print(target_machines) + return target_machines[0] + except UnknownRecordError: + machine = Machine( + id=self._machine_repository.get_new_id(), + network_interfaces=[IPv4Interface(event.target)], + ) + self._machine_repository.upsert_machine(machine) + return machine + + def _update_nodes(self, target_machine: Machine, event: ScanEvent): + src_machine = self._get_source_machine(event) + + self._node_repository.upsert_communication( + src_machine.id, target_machine.id, CommunicationType.SCANNED + ) + + 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) + + def _update_target_machine_os(self, machine: Machine, event: PingScanEvent): + if event.os is not None and machine.operating_system is None: + machine.operating_system = event.os + self._machine_repository.upsert_machine(machine) From 2686a7a4ee5fd49c3226b265eb9989cfd3dec68b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:16:31 +0200 Subject: [PATCH 02/24] UT: Add tests for ScanEventHandler --- .../test_scan_event_handler.py | 336 ++++++++++++++++++ 1 file changed, 336 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py 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 new file mode 100644 index 000000000..2a6d11b0a --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_scan_event_handler.py @@ -0,0 +1,336 @@ +from ipaddress import IPv4Address, IPv4Interface +from itertools import count +from unittest.mock import MagicMock +from uuid import UUID + +import pytest + +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.repository import ( + IAgentRepository, + IMachineRepository, + INodeRepository, + RetrievalError, + StorageError, + UnknownRecordError, +) + +SEED_ID = 99 +AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2") +MACHINE_ID = 11 +CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000") +AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER) +SOURCE_MACHINE = Machine( + id=MACHINE_ID, + hardware_id=5, + network_interfaces=[IPv4Interface("10.10.10.99/24")], +) +TARGET_MACHINE = Machine( + id=33, + hardware_id=9, + network_interfaces=[IPv4Interface("10.10.10.1/24")], +) + +PING_SCAN_EVENT = PingScanEvent( + source=AGENT_ID, + target=IPv4Address("10.10.10.1"), + response_received=True, + os=OperatingSystem.LINUX, +) + +PING_SCAN_EVENT_NO_RESPONSE = PingScanEvent( + source=AGENT_ID, + target=IPv4Address("10.10.10.1"), + response_received=False, + os=OperatingSystem.LINUX, +) + +PING_SCAN_EVENT_NO_OS = PingScanEvent( + source=AGENT_ID, + target=IPv4Address("10.10.10.1"), + 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}, +) + +TCP_SCAN_EVENT_CLOSED = TCPScanEvent( + source=AGENT_ID, + target=IPv4Address("10.10.10.1"), + ports={145: PortStatus.CLOSED, 8080: PortStatus.CLOSED}, +) + + +@pytest.fixture +def agent_repository() -> IAgentRepository: + agent_repository = MagicMock(spec=IAgentRepository) + agent_repository.upsert_agent = MagicMock() + agent_repository.get_agent_by_id = MagicMock(return_value=AGENT) + return agent_repository + + +@pytest.fixture +def machine_repository() -> IMachineRepository: + machine_repository = MagicMock(spec=IMachineRepository) + machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID)) + machine_repository.upsert_machine = MagicMock() + return machine_repository + + +@pytest.fixture +def node_repository() -> INodeRepository: + node_repository = MagicMock(spec=INodeRepository) + node_repository.upsert_communication = MagicMock() + return node_repository + + +@pytest.fixture +def scan_event_handler(agent_repository, machine_repository, node_repository): + return ScanEventHandler(agent_repository, machine_repository, node_repository) + + +@pytest.fixture +def handle_ping_scan_event(scan_event_handler): + return scan_event_handler.handle_ping_scan_event + + +@pytest.fixture +def handle_tcp_scan_event(scan_event_handler): + return scan_event_handler.handle_tcp_scan_event + + +machines = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} +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], +} + + +@pytest.fixture(params=[SOURCE_MACHINE.id, TARGET_MACHINE.id]) +def machine_id(request): + return request.param + + +def machine_from_id(id: int): + return machines_by_id[id] + + +def machines_from_ip(ip: IPv4Address): + return machines_by_ip[ip] + + +class error_machine_by_id: + """Raise an error if the machine with the called ID matches the stored ID""" + + def __init__(self, id: int, error): + self.id = id + self.error = error + + def __call__(self, id: int): + if id == self.id: + raise self.error + else: + return machine_from_id(id) + + +class error_machine_by_ip: + """Raise an error if the machine with the called IP matches the stored ID""" + + def __init__(self, id: int, error): + self.id = id + self.error = error + + def __call__(self, ip: IPv4Address): + print(f"IP is: {ip}") + machines = machines_from_ip(ip) + if machines[0].id == self.id: + print(f"Raise error: {self.error}") + raise self.error + else: + print(f"Return machine: {machines}") + return machines + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__target_machine_not_exists( + event, handler, machine_repository: IMachineRepository, request +): + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + + handler = request.getfixturevalue(handler) + handler(event) + + expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(event.target)]) + if event == PING_SCAN_EVENT: + expected_machine.operating_system = event.os + machine_repository.upsert_machine.assert_called_with(expected_machine) + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__upserts_node( + event, + handler, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + request, +): + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) + + handler = request.getfixturevalue(handler) + handler(event) + + node_repository.upsert_communication.assert_called_with( + SOURCE_MACHINE.id, TARGET_MACHINE.id, CommunicationType.SCANNED + ) + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__node_not_upserted_if_no_matching_agent( + event, + handler, + agent_repository: IAgentRepository, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + request, +): + agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError) + machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not node_repository.upsert_communication.called + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( + event, + handler, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + request, + machine_id, +): + machine_repository.get_machine_by_id = MagicMock( + side_effect=error_machine_by_id(machine_id, RetrievalError) + ) + machine_repository.get_machines_by_ip = MagicMock( + side_effect=error_machine_by_ip(machine_id, RetrievalError) + ) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not node_repository.upsert_communication.called + + +@pytest.mark.parametrize( + "event,handler", + [ + (PING_SCAN_EVENT_NO_OS, "handle_ping_scan_event"), + (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), + ], +) +def test_scan_event_handler__machine_not_upserted( + event, handler, machine_repository: IMachineRepository, request +): + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not machine_repository.upsert_machine.called + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( + event, handler, machine_repository: IMachineRepository, request +): + machine_with_os = TARGET_MACHINE + machine_with_os.operating_system = OperatingSystem.WINDOWS + machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not machine_repository.upsert_machine.called + + +@pytest.mark.parametrize( + "event,handler", + [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], +) +def test_scan_event_handler__node_not_upserted_if_machine_storageerror( + event, + handler, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + request, +): + if event == PING_SCAN_EVENT: + target_machine = TARGET_MACHINE + target_machine.operating_system = None + + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + if event == PING_SCAN_EVENT: + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) + machine_repository.upsert_machine = MagicMock(side_effect=StorageError) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not node_repository.upsert_communication.called + + +@pytest.mark.parametrize( + "event,handler", + [ + (PING_SCAN_EVENT_NO_RESPONSE, "handle_ping_scan_event"), + (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), + ], +) +def test_scan_event_handler__failed_scan( + event, + handler, + machine_repository: IMachineRepository, + node_repository: INodeRepository, + request, +): + machine_repository.upsert_machine = MagicMock(side_effect=StorageError) + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) + + handler = request.getfixturevalue(handler) + handler(event) + + assert not node_repository.upsert_communication.called + assert not machine_repository.upsert_machine.called From 8002080c8b49630a7f52c7e106fdd9f478a3148b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:17:37 +0200 Subject: [PATCH 03/24] Island: Subscribe to Ping and TCP scan events with ScanEventHandlers --- .../cc/agent_event_handlers/__init__.py | 1 + .../monkey_island/cc/setup/agent_event_handlers.py | 13 +++++++------ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/agent_event_handlers/__init__.py b/monkey/monkey_island/cc/agent_event_handlers/__init__.py index 1c7e4714c..956bc6717 100644 --- a/monkey/monkey_island/cc/agent_event_handlers/__init__.py +++ b/monkey/monkey_island/cc/agent_event_handlers/__init__.py @@ -1,3 +1,4 @@ from .handle_ping_scan_event import handle_ping_scan_event from .save_event_to_event_repository import save_event_to_event_repository from .save_stolen_credentials_to_repository import save_stolen_credentials_to_repository +from .scan_event_handler import ScanEventHandler diff --git a/monkey/monkey_island/cc/setup/agent_event_handlers.py b/monkey/monkey_island/cc/setup/agent_event_handlers.py index 2cc1bf3da..2981349de 100644 --- a/monkey/monkey_island/cc/setup/agent_event_handlers.py +++ b/monkey/monkey_island/cc/setup/agent_event_handlers.py @@ -1,8 +1,8 @@ from common import DIContainer -from common.agent_events import CredentialsStolenEvent, PingScanEvent +from common.agent_events import CredentialsStolenEvent, PingScanEvent, TCPScanEvent from common.event_queue import IAgentEventQueue from monkey_island.cc.agent_event_handlers import ( - handle_ping_scan_event, + ScanEventHandler, save_event_to_event_repository, save_stolen_credentials_to_repository, ) @@ -17,7 +17,7 @@ from monkey_island.cc.repository import ( def setup_agent_event_handlers(container: DIContainer): _subscribe_and_store_to_event_repository(container) - _subscribe_ping_scan_event(container) + _subscribe_scan_events(container) def _subscribe_and_store_to_event_repository(container: DIContainer): @@ -32,12 +32,13 @@ def _subscribe_and_store_to_event_repository(container: DIContainer): agent_event_queue.subscribe_type(CredentialsStolenEvent, save_stolen_credentials_subscriber) -def _subscribe_ping_scan_event(container: DIContainer): +def _subscribe_scan_events(container: DIContainer): agent_event_queue = container.resolve(IAgentEventQueue) agent_repository = container.resolve(IAgentRepository) machine_repository = container.resolve(IMachineRepository) node_repository = container.resolve(INodeRepository) - handler = handle_ping_scan_event(agent_repository, machine_repository, node_repository) + scan_event_handler = ScanEventHandler(agent_repository, machine_repository, node_repository) - agent_event_queue.subscribe_type(PingScanEvent, handler) + agent_event_queue.subscribe_type(PingScanEvent, scan_event_handler.handle_ping_scan_event) + agent_event_queue.subscribe_type(TCPScanEvent, scan_event_handler.handle_tcp_scan_event) From f23093dc783291449f421e839d75fc91bf86aa79 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:20:23 +0200 Subject: [PATCH 04/24] Island: Remove handle_ping_scan_event callable class --- .../cc/agent_event_handlers/__init__.py | 1 - .../handle_ping_scan_event.py | 76 ----- .../test_handle_ping_scan_event.py | 268 ------------------ 3 files changed, 345 deletions(-) delete mode 100644 monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py delete mode 100644 monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py diff --git a/monkey/monkey_island/cc/agent_event_handlers/__init__.py b/monkey/monkey_island/cc/agent_event_handlers/__init__.py index 956bc6717..b04d81961 100644 --- a/monkey/monkey_island/cc/agent_event_handlers/__init__.py +++ b/monkey/monkey_island/cc/agent_event_handlers/__init__.py @@ -1,4 +1,3 @@ -from .handle_ping_scan_event import handle_ping_scan_event from .save_event_to_event_repository import save_event_to_event_repository from .save_stolen_credentials_to_repository import save_stolen_credentials_to_repository from .scan_event_handler import ScanEventHandler diff --git a/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py b/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py deleted file mode 100644 index 3e5d4041f..000000000 --- a/monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py +++ /dev/null @@ -1,76 +0,0 @@ -from logging import getLogger - -from pydantic.json import IPv4Interface - -from common.agent_events import PingScanEvent -from monkey_island.cc.models import CommunicationType, Machine -from monkey_island.cc.repository import ( - IAgentRepository, - IMachineRepository, - INodeRepository, - RetrievalError, - StorageError, - UnknownRecordError, -) - -logger = getLogger(__name__) - - -class handle_ping_scan_event: - """ - Handles ping scan event and makes changes to Machine and Node states based on it - """ - - def __init__( - self, - agent_repository: IAgentRepository, - machine_repository: IMachineRepository, - node_repository: INodeRepository, - ): - self._agent_repository = agent_repository - self._machine_repository = machine_repository - self._node_repository = node_repository - - def __call__(self, event: PingScanEvent): - """ - :param event: Ping scan event to process - """ - if not event.response_received: - return - - try: - target_machine = self._get_target_machine(event) - - self._update_target_machine_os(target_machine, event) - self._update_nodes(target_machine, event) - except (RetrievalError, StorageError, UnknownRecordError): - logger.exception("Unable to process ping scan data") - - def _get_target_machine(self, event: PingScanEvent) -> Machine: - try: - target_machines = self._machine_repository.get_machines_by_ip(event.target) - return target_machines[0] - except UnknownRecordError: - machine = Machine( - id=self._machine_repository.get_new_id(), - network_interfaces=[IPv4Interface(event.target)], - ) - self._machine_repository.upsert_machine(machine) - return machine - - def _update_target_machine_os(self, machine: Machine, event: PingScanEvent): - if event.os is not None and machine.operating_system is None: - machine.operating_system = event.os - self._machine_repository.upsert_machine(machine) - - def _update_nodes(self, target_machine, event): - src_machine = self._get_source_machine(event) - - # Update or create the node - self._node_repository.upsert_communication( - src_machine.id, target_machine.id, CommunicationType.SCANNED - ) - - def _get_source_machine(self, event: PingScanEvent) -> 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_handle_ping_scan_event.py b/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py deleted file mode 100644 index 9849e893e..000000000 --- a/monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py +++ /dev/null @@ -1,268 +0,0 @@ -from ipaddress import IPv4Address, IPv4Interface -from itertools import count -from unittest.mock import MagicMock -from uuid import UUID - -import pytest - -from common import OperatingSystem -from common.agent_events import PingScanEvent -from common.types import SocketAddress -from monkey_island.cc.agent_event_handlers import handle_ping_scan_event -from monkey_island.cc.models import Agent, CommunicationType, Machine -from monkey_island.cc.repository import ( - IAgentRepository, - IMachineRepository, - INodeRepository, - RetrievalError, - StorageError, - UnknownRecordError, -) - -SEED_ID = 99 -AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2") -MACHINE_ID = 11 -CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000") -AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER) -PINGER_MACHINE = Machine( - id=MACHINE_ID, - hardware_id=5, - network_interfaces=[IPv4Interface("10.10.10.99/24")], -) -TARGET_MACHINE = Machine( - id=33, - hardware_id=9, - network_interfaces=[IPv4Interface("10.10.10.1/24")], -) -EVENT = PingScanEvent( - source=AGENT_ID, - target=IPv4Address("10.10.10.1"), - response_received=True, - os=OperatingSystem.LINUX, -) - -EVENT_NO_RESPONSE = PingScanEvent( - source=AGENT_ID, - target=IPv4Address("10.10.10.1"), - response_received=False, - os=OperatingSystem.LINUX, -) - - -@pytest.fixture -def agent_repository() -> IAgentRepository: - agent_repository = MagicMock(spec=IAgentRepository) - agent_repository.upsert_agent = MagicMock() - agent_repository.get_agent_by_id = MagicMock(return_value=AGENT) - return agent_repository - - -@pytest.fixture -def machine_repository() -> IMachineRepository: - machine_repository = MagicMock(spec=IMachineRepository) - machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID)) - machine_repository.upsert_machine = MagicMock() - return machine_repository - - -@pytest.fixture -def node_repository() -> INodeRepository: - node_repository = MagicMock(spec=INodeRepository) - node_repository.upsert_communication = MagicMock() - return node_repository - - -@pytest.fixture -def handler(agent_repository, machine_repository, node_repository) -> handle_ping_scan_event: - return handle_ping_scan_event(agent_repository, machine_repository, node_repository) - - -machines = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} -machines_by_id = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} -machines_by_ip = { - IPv4Address("10.10.10.99"): [PINGER_MACHINE], - IPv4Address("10.10.10.1"): [TARGET_MACHINE], -} - - -def machine_from_id(id: int): - return machines_by_id[id] - - -def machines_from_ip(ip: IPv4Address): - return machines_by_ip[ip] - - -class error_machine_by_id: - """Raise an error if the machine with the called ID matches the stored ID""" - - def __init__(self, id: int, error): - self.id = id - self.error = error - - def __call__(self, id: int): - if id == self.id: - raise self.error - else: - return machine_from_id(id) - - -class error_machine_by_ip: - """Raise an error if the machine with the called IP matches the stored ID""" - - def __init__(self, id: int, error): - self.id = id - self.error = error - - def __call__(self, ip: IPv4Address): - print(f"IP is: {ip}") - machines = machines_from_ip(ip) - if machines[0].id == self.id: - print(f"Raise error: {self.error}") - raise self.error - else: - print(f"Return machine: {machines}") - return machines - - -def test_handle_ping_scan_event__target_machine_not_exists( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, -): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) - - handler(EVENT) - - expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(EVENT.target)]) - expected_machine.operating_system = EVENT.os - machine_repository.upsert_machine.assert_called_with(expected_machine) - - -def test_handle_ping_scan_event__target_machine_already_exists( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, -): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - - handler(EVENT) - - expected_machine = TARGET_MACHINE.copy() - expected_machine.operating_system = OperatingSystem.LINUX - machine_repository.upsert_machine.assert_called_with(expected_machine) - - -def test_handle_ping_scan_event__upserts_node( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, - node_repository: INodeRepository, -): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) - - handler(EVENT) - - node_repository.upsert_communication.assert_called_with( - PINGER_MACHINE.id, TARGET_MACHINE.id, CommunicationType.SCANNED - ) - - -def test_handle_ping_scan_event__node_not_upserted_if_no_matching_agent( - handler: handle_ping_scan_event, - agent_repository: IAgentRepository, - machine_repository: IMachineRepository, - node_repository: INodeRepository, -): - agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError) - machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE) - - handler(EVENT) - - assert not node_repository.upsert_communication.called - - -def test_handle_ping_scan_event__node_not_upserted_if_no_matching_machine( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, - node_repository: INodeRepository, -): - machine_repository.get_machine_by_id = MagicMock(side_effect=UnknownRecordError) - - handler(EVENT) - - assert not node_repository.upsert_communication.called - - -@pytest.mark.parametrize("id", [PINGER_MACHINE.id, TARGET_MACHINE.id]) -def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, - node_repository: INodeRepository, - id, -): - machine_repository.get_machine_by_id = MagicMock( - side_effect=error_machine_by_id(id, RetrievalError) - ) - machine_repository.get_machines_by_ip = MagicMock( - side_effect=error_machine_by_ip(id, RetrievalError) - ) - - handler(EVENT) - - assert not node_repository.upsert_communication.called - - -def test_handle_scan_data__machine_not_upserted_if_os_is_none( - handler: handle_ping_scan_event, machine_repository: IMachineRepository -): - event = PingScanEvent(source=EVENT.source, target=EVENT.target, response_received=True, os=None) - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - - handler(event) - - assert not machine_repository.upsert_machine.called - - -def test_handle_scan_data__machine_not_upserted_if_existing_machine_has_os( - handler: handle_ping_scan_event, machine_repository: IMachineRepository -): - machine_with_os = TARGET_MACHINE - machine_with_os.operating_system = OperatingSystem.WINDOWS - machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os) - - handler(EVENT) - - assert not machine_repository.upsert_machine.called - - -def test_handle_scan_data__node_not_upserted_if_machine_storageerror( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, - node_repository: INodeRepository, -): - target_machine = TARGET_MACHINE - target_machine.operating_system = None - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - machine_repository.upsert_machine = MagicMock(side_effect=StorageError) - - handler(EVENT) - - assert not node_repository.upsert_communication.called - - -def test_handle_scan_data__failed_ping( - handler: handle_ping_scan_event, - machine_repository: IMachineRepository, - node_repository: INodeRepository, -): - machine_repository.upsert_machine = MagicMock(side_effect=StorageError) - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - - handler(EVENT_NO_RESPONSE) - - assert not node_repository.upsert_communication.called - assert not machine_repository.upsert_machine.called From d42a353aaa593747ab889608fbfea2e2fbdca2ec Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:23:04 +0200 Subject: [PATCH 05/24] Island: Remove TypeError exception from ScanEventHandler --- .../cc/agent_event_handlers/scan_event_handler.py | 3 +-- 1 file changed, 1 insertion(+), 2 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 2d1da09e6..e5c36b3cc 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 @@ -45,7 +45,7 @@ class ScanEventHandler: self._update_target_machine_os(target_machine, event) self._update_nodes(target_machine, event) - except (RetrievalError, StorageError, TypeError, UnknownRecordError): + except (RetrievalError, StorageError, UnknownRecordError): logger.exception("Unable to process ping scan data") def handle_tcp_scan_event(self, event: TCPScanEvent): @@ -64,7 +64,6 @@ class ScanEventHandler: def _get_target_machine(self, event: ScanEvent) -> Machine: try: target_machines = self._machine_repository.get_machines_by_ip(event.target) - print(target_machines) return target_machines[0] except UnknownRecordError: machine = Machine( From 326b07e5c1165a75513ed977bfde007ad6e10cd5 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 30 Sep 2022 16:26:00 +0200 Subject: [PATCH 06/24] Island: Fix typo in ScanEventHandler docstring --- .../monkey_island/cc/agent_event_handlers/scan_event_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e5c36b3cc..6e3582832 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 @@ -23,7 +23,7 @@ logger = getLogger(__name__) class ScanEventHandler: """ - Handles scan event and makes changes to Machine and Node states base on it + Handles scan event and makes changes to Machine and Node states based on it """ def __init__( From f6e5462ad31f2d77623bdb6ef8ef0984555f9d17 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:06:09 -0400 Subject: [PATCH 07/24] Island: Use sum() to count open ports --- .../cc/agent_event_handlers/scan_event_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 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 6e3582832..673074591 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 @@ -49,9 +49,9 @@ class ScanEventHandler: logger.exception("Unable to process ping scan data") def handle_tcp_scan_event(self, event: TCPScanEvent): - open_ports = [port for port, status in event.ports.items() if status == PortStatus.OPEN] + num_open_ports = sum((1 for status in event.ports.values() if status == PortStatus.OPEN)) - if not open_ports: + if num_open_ports <= 0: return try: From 0c786dfd94b0adda99fb7fa294ca3a6f2e101015 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:08:32 -0400 Subject: [PATCH 08/24] Island: Change method order in ScanEventHandler --- .../cc/agent_event_handlers/scan_event_handler.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 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 673074591..0eac3637c 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 @@ -73,6 +73,11 @@ class ScanEventHandler: self._machine_repository.upsert_machine(machine) return machine + def _update_target_machine_os(self, machine: Machine, event: PingScanEvent): + if event.os is not None and machine.operating_system is None: + machine.operating_system = event.os + self._machine_repository.upsert_machine(machine) + def _update_nodes(self, target_machine: Machine, event: ScanEvent): src_machine = self._get_source_machine(event) @@ -83,8 +88,3 @@ class ScanEventHandler: 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) - - def _update_target_machine_os(self, machine: Machine, event: PingScanEvent): - if event.os is not None and machine.operating_system is None: - machine.operating_system = event.os - self._machine_repository.upsert_machine(machine) From 07fa283ce1712e4182e500d713fc36d8feb5ca50 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:28:22 -0400 Subject: [PATCH 09/24] Island: Use container.resolve(ScanEventHandler) --- .../cc/setup/agent_event_handlers.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/monkey/monkey_island/cc/setup/agent_event_handlers.py b/monkey/monkey_island/cc/setup/agent_event_handlers.py index 2981349de..572e6de3d 100644 --- a/monkey/monkey_island/cc/setup/agent_event_handlers.py +++ b/monkey/monkey_island/cc/setup/agent_event_handlers.py @@ -6,13 +6,7 @@ from monkey_island.cc.agent_event_handlers import ( save_event_to_event_repository, save_stolen_credentials_to_repository, ) -from monkey_island.cc.repository import ( - IAgentEventRepository, - IAgentRepository, - ICredentialsRepository, - IMachineRepository, - INodeRepository, -) +from monkey_island.cc.repository import IAgentEventRepository, ICredentialsRepository def setup_agent_event_handlers(container: DIContainer): @@ -23,6 +17,7 @@ def setup_agent_event_handlers(container: DIContainer): def _subscribe_and_store_to_event_repository(container: DIContainer): agent_event_queue = container.resolve(IAgentEventQueue) + # TODO: Can't we just `container.resolve(save_event_to_event_repository)`? save_event_subscriber = save_event_to_event_repository(container.resolve(IAgentEventRepository)) agent_event_queue.subscribe_all_events(save_event_subscriber) @@ -34,11 +29,7 @@ def _subscribe_and_store_to_event_repository(container: DIContainer): def _subscribe_scan_events(container: DIContainer): agent_event_queue = container.resolve(IAgentEventQueue) - agent_repository = container.resolve(IAgentRepository) - machine_repository = container.resolve(IMachineRepository) - node_repository = container.resolve(INodeRepository) - - scan_event_handler = ScanEventHandler(agent_repository, machine_repository, node_repository) + scan_event_handler = container.resolve(ScanEventHandler) agent_event_queue.subscribe_type(PingScanEvent, scan_event_handler.handle_ping_scan_event) agent_event_queue.subscribe_type(TCPScanEvent, scan_event_handler.handle_tcp_scan_event) From e2c86d3d7a78bee0e793ce1e453b280278531eac Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Fri, 30 Sep 2022 10:33:59 +0000 Subject: [PATCH 10/24] Common: Register PingScanEvent to event serializer --- monkey/common/agent_event_serializers/register.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/common/agent_event_serializers/register.py b/monkey/common/agent_event_serializers/register.py index fbe353db9..0a820f714 100644 --- a/monkey/common/agent_event_serializers/register.py +++ b/monkey/common/agent_event_serializers/register.py @@ -1,4 +1,4 @@ -from common.agent_events import CredentialsStolenEvent +from common.agent_events import CredentialsStolenEvent, PingScanEvent from . import AgentEventSerializerRegistry, PydanticAgentEventSerializer @@ -9,3 +9,4 @@ def register_common_agent_event_serializers( event_serializer_registry[CredentialsStolenEvent] = PydanticAgentEventSerializer( CredentialsStolenEvent ) + event_serializer_registry[PingScanEvent] = PydanticAgentEventSerializer(PingScanEvent) From f89068ae0044a3f94363deda311fe39504654632 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:35:57 -0400 Subject: [PATCH 11/24] Common: Register TCPScanEvent serializer --- monkey/common/agent_event_serializers/register.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/common/agent_event_serializers/register.py b/monkey/common/agent_event_serializers/register.py index 0a820f714..1ccf6cadf 100644 --- a/monkey/common/agent_event_serializers/register.py +++ b/monkey/common/agent_event_serializers/register.py @@ -1,4 +1,4 @@ -from common.agent_events import CredentialsStolenEvent, PingScanEvent +from common.agent_events import CredentialsStolenEvent, PingScanEvent, TCPScanEvent from . import AgentEventSerializerRegistry, PydanticAgentEventSerializer @@ -10,3 +10,4 @@ def register_common_agent_event_serializers( CredentialsStolenEvent ) event_serializer_registry[PingScanEvent] = PydanticAgentEventSerializer(PingScanEvent) + event_serializer_registry[TCPScanEvent] = PydanticAgentEventSerializer(TCPScanEvent) From 20b84aa1a418274ce2d2942f20a49e673b8a8234 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:48:03 -0400 Subject: [PATCH 12/24] UT: Use indirect instead of get_fixture_value() --- .../test_scan_event_handler.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 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 2a6d11b0a..d712d854f 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 @@ -161,9 +161,15 @@ class error_machine_by_ip: return machines +@pytest.fixture +def handler(scan_event_handler, request): + return getattr(scan_event_handler, request.param) + + @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__target_machine_not_exists( event, handler, machine_repository: IMachineRepository, request @@ -171,7 +177,6 @@ def test_scan_event_handler__target_machine_not_exists( machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) - handler = request.getfixturevalue(handler) handler(event) expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(event.target)]) @@ -183,6 +188,7 @@ def test_scan_event_handler__target_machine_not_exists( @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__upserts_node( event, @@ -194,7 +200,6 @@ def test_scan_event_handler__upserts_node( machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) - handler = request.getfixturevalue(handler) handler(event) node_repository.upsert_communication.assert_called_with( @@ -205,6 +210,7 @@ def test_scan_event_handler__upserts_node( @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_no_matching_agent( event, @@ -217,7 +223,6 @@ def test_scan_event_handler__node_not_upserted_if_no_matching_agent( agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError) machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE) - handler = request.getfixturevalue(handler) handler(event) assert not node_repository.upsert_communication.called @@ -226,6 +231,7 @@ def test_scan_event_handler__node_not_upserted_if_no_matching_agent( @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( event, @@ -242,7 +248,6 @@ def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( side_effect=error_machine_by_ip(machine_id, RetrievalError) ) - handler = request.getfixturevalue(handler) handler(event) assert not node_repository.upsert_communication.called @@ -254,6 +259,7 @@ def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( (PING_SCAN_EVENT_NO_OS, "handle_ping_scan_event"), (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), ], + indirect=["handler"], ) def test_scan_event_handler__machine_not_upserted( event, handler, machine_repository: IMachineRepository, request @@ -261,7 +267,6 @@ def test_scan_event_handler__machine_not_upserted( machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - handler = request.getfixturevalue(handler) handler(event) assert not machine_repository.upsert_machine.called @@ -270,6 +275,7 @@ def test_scan_event_handler__machine_not_upserted( @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( event, handler, machine_repository: IMachineRepository, request @@ -278,7 +284,6 @@ def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( machine_with_os.operating_system = OperatingSystem.WINDOWS machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os) - handler = request.getfixturevalue(handler) handler(event) assert not machine_repository.upsert_machine.called @@ -287,6 +292,7 @@ def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( @pytest.mark.parametrize( "event,handler", [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_machine_storageerror( event, @@ -305,7 +311,6 @@ def test_scan_event_handler__node_not_upserted_if_machine_storageerror( machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) machine_repository.upsert_machine = MagicMock(side_effect=StorageError) - handler = request.getfixturevalue(handler) handler(event) assert not node_repository.upsert_communication.called @@ -317,6 +322,7 @@ def test_scan_event_handler__node_not_upserted_if_machine_storageerror( (PING_SCAN_EVENT_NO_RESPONSE, "handle_ping_scan_event"), (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), ], + indirect=["handler"], ) def test_scan_event_handler__failed_scan( event, @@ -329,7 +335,6 @@ def test_scan_event_handler__failed_scan( machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - handler = request.getfixturevalue(handler) handler(event) assert not node_repository.upsert_communication.called From 79f56e0789abee710ae8dc6b4b28cc598048bf88 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:53:47 -0400 Subject: [PATCH 13/24] UT: Use constants for TCP and ping scan method names --- .../test_scan_event_handler.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 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 d712d854f..a14607197 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 @@ -161,6 +161,10 @@ class error_machine_by_ip: return machines +HANDLE_PING_SCAN_METHOD = "handle_ping_scan_event" +HANDLE_TCP_SCAN_METHOD = "handle_tcp_scan_event" + + @pytest.fixture def handler(scan_event_handler, request): return getattr(scan_event_handler, request.param) @@ -168,7 +172,7 @@ def handler(scan_event_handler, request): @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__target_machine_not_exists( @@ -187,7 +191,7 @@ def test_scan_event_handler__target_machine_not_exists( @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__upserts_node( @@ -209,7 +213,7 @@ def test_scan_event_handler__upserts_node( @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_no_matching_agent( @@ -230,7 +234,7 @@ def test_scan_event_handler__node_not_upserted_if_no_matching_agent( @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( @@ -256,8 +260,8 @@ def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( @pytest.mark.parametrize( "event,handler", [ - (PING_SCAN_EVENT_NO_OS, "handle_ping_scan_event"), - (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), + (PING_SCAN_EVENT_NO_OS, HANDLE_PING_SCAN_METHOD), + (TCP_SCAN_EVENT_CLOSED, HANDLE_TCP_SCAN_METHOD), ], indirect=["handler"], ) @@ -274,7 +278,7 @@ def test_scan_event_handler__machine_not_upserted( @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( @@ -291,7 +295,7 @@ def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( @pytest.mark.parametrize( "event,handler", - [(PING_SCAN_EVENT, "handle_ping_scan_event"), (TCP_SCAN_EVENT, "handle_tcp_scan_event")], + [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) def test_scan_event_handler__node_not_upserted_if_machine_storageerror( @@ -319,8 +323,8 @@ def test_scan_event_handler__node_not_upserted_if_machine_storageerror( @pytest.mark.parametrize( "event,handler", [ - (PING_SCAN_EVENT_NO_RESPONSE, "handle_ping_scan_event"), - (TCP_SCAN_EVENT_CLOSED, "handle_tcp_scan_event"), + (PING_SCAN_EVENT_NO_RESPONSE, HANDLE_PING_SCAN_METHOD), + (TCP_SCAN_EVENT_CLOSED, HANDLE_TCP_SCAN_METHOD), ], indirect=["handler"], ) From 029c1016435f0d464f1c0b960e67deb02d6b8161 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 12:54:27 -0400 Subject: [PATCH 14/24] UT: Remove disused scan event fixtures --- .../cc/agent_event_handlers/test_scan_event_handler.py | 10 ---------- 1 file changed, 10 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 a14607197..b22d86c16 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 @@ -98,16 +98,6 @@ def scan_event_handler(agent_repository, machine_repository, node_repository): return ScanEventHandler(agent_repository, machine_repository, node_repository) -@pytest.fixture -def handle_ping_scan_event(scan_event_handler): - return scan_event_handler.handle_ping_scan_event - - -@pytest.fixture -def handle_tcp_scan_event(scan_event_handler): - return scan_event_handler.handle_tcp_scan_event - - machines = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} machines_by_id = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} machines_by_ip = { From 97061ea61ca05ea005a37c710e75111b7a28e2af Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 13:04:57 -0400 Subject: [PATCH 15/24] UT: Remove redundant "test_scan_event_handler" from test names --- .../test_scan_event_handler.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 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 b22d86c16..b716ca32f 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 @@ -165,9 +165,7 @@ def handler(scan_event_handler, request): [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__target_machine_not_exists( - event, handler, machine_repository: IMachineRepository, request -): +def test_target_machine_not_exists(event, handler, machine_repository: IMachineRepository, request): machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) @@ -184,7 +182,7 @@ def test_scan_event_handler__target_machine_not_exists( [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__upserts_node( +def test_upserts_node( event, handler, machine_repository: IMachineRepository, @@ -206,7 +204,7 @@ def test_scan_event_handler__upserts_node( [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__node_not_upserted_if_no_matching_agent( +def test_node_not_upserted_if_no_matching_agent( event, handler, agent_repository: IAgentRepository, @@ -227,7 +225,7 @@ def test_scan_event_handler__node_not_upserted_if_no_matching_agent( [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( +def test_node_not_upserted_if_machine_retrievalerror( event, handler, machine_repository: IMachineRepository, @@ -255,9 +253,7 @@ def test_scan_event_handler__node_not_upserted_if_machine_retrievalerror( ], indirect=["handler"], ) -def test_scan_event_handler__machine_not_upserted( - event, handler, machine_repository: IMachineRepository, request -): +def test_machine_not_upserted(event, handler, machine_repository: IMachineRepository, request): machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) @@ -271,7 +267,7 @@ def test_scan_event_handler__machine_not_upserted( [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( +def test_machine_not_upserted_if_existing_machine_has_os( event, handler, machine_repository: IMachineRepository, request ): machine_with_os = TARGET_MACHINE @@ -288,7 +284,7 @@ def test_scan_event_handler__machine_not_upserted_if_existing_machine_has_os( [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_scan_event_handler__node_not_upserted_if_machine_storageerror( +def test_node_not_upserted_if_machine_storageerror( event, handler, machine_repository: IMachineRepository, @@ -318,7 +314,7 @@ def test_scan_event_handler__node_not_upserted_if_machine_storageerror( ], indirect=["handler"], ) -def test_scan_event_handler__failed_scan( +def test_failed_scan( event, handler, machine_repository: IMachineRepository, From 281100901987b8b431b3fb6cb9c50ecd9256af87 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 13:10:57 -0400 Subject: [PATCH 16/24] UT: Deparametrize a test TCP scan events do not contain OS information. --- .../cc/agent_event_handlers/test_scan_event_handler.py | 9 ++------- 1 file changed, 2 insertions(+), 7 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 b716ca32f..c2325ed45 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 @@ -262,19 +262,14 @@ def test_machine_not_upserted(event, handler, machine_repository: IMachineReposi assert not machine_repository.upsert_machine.called -@pytest.mark.parametrize( - "event,handler", - [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], - indirect=["handler"], -) def test_machine_not_upserted_if_existing_machine_has_os( - event, handler, machine_repository: IMachineRepository, request + scan_event_handler, machine_repository: IMachineRepository, request ): machine_with_os = TARGET_MACHINE machine_with_os.operating_system = OperatingSystem.WINDOWS machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os) - handler(event) + scan_event_handler.handle_ping_scan_event(PING_SCAN_EVENT) assert not machine_repository.upsert_machine.called From 067a143f2c6810cab2522c66e127ea263adc5649 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 13:16:04 -0400 Subject: [PATCH 17/24] UT: Fix bug in test mock --- .../cc/agent_event_handlers/test_scan_event_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c2325ed45..82ede960f 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 @@ -267,7 +267,7 @@ def test_machine_not_upserted_if_existing_machine_has_os( ): machine_with_os = TARGET_MACHINE machine_with_os.operating_system = OperatingSystem.WINDOWS - machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os) + machine_repository.get_machines_by_ip = MagicMock(return_value=[machine_with_os]) scan_event_handler.handle_ping_scan_event(PING_SCAN_EVENT) From 1dbfca567a0d674938d49588e7de3f456f8a719b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:28:26 +0000 Subject: [PATCH 18/24] UT: Give constants uppercase names --- .../cc/agent_event_handlers/test_scan_event_handler.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 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 82ede960f..ae8f7f2d9 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 @@ -98,9 +98,9 @@ def scan_event_handler(agent_repository, machine_repository, node_repository): return ScanEventHandler(agent_repository, machine_repository, node_repository) -machines = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} -machines_by_id = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} -machines_by_ip = { +MACHINES = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} +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], } @@ -112,11 +112,11 @@ def machine_id(request): def machine_from_id(id: int): - return machines_by_id[id] + return MACHINES_BY_ID[id] def machines_from_ip(ip: IPv4Address): - return machines_by_ip[ip] + return MACHINES_BY_IP[ip] class error_machine_by_id: From cf13481865fc06e77dedff313998717a1bae388a Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:29:27 +0000 Subject: [PATCH 19/24] UT: Remove unused constant --- .../cc/agent_event_handlers/test_scan_event_handler.py | 1 - 1 file changed, 1 deletion(-) 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 ae8f7f2d9..c6a71183b 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 @@ -98,7 +98,6 @@ def scan_event_handler(agent_repository, machine_repository, node_repository): return ScanEventHandler(agent_repository, machine_repository, node_repository) -MACHINES = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} MACHINES_BY_ID = {MACHINE_ID: SOURCE_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE} MACHINES_BY_IP = { IPv4Address("10.10.10.99"): [SOURCE_MACHINE], From da5d7b7357fb7925cb40d1596542637453c6e4c6 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:32:08 +0000 Subject: [PATCH 20/24] UT: Removed unsed 'request' parameter --- .../cc/agent_event_handlers/test_scan_event_handler.py | 5 ----- 1 file changed, 5 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 c6a71183b..5204a30ea 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 @@ -186,7 +186,6 @@ def test_upserts_node( handler, machine_repository: IMachineRepository, node_repository: INodeRepository, - request, ): machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) @@ -209,7 +208,6 @@ def test_node_not_upserted_if_no_matching_agent( agent_repository: IAgentRepository, machine_repository: IMachineRepository, node_repository: INodeRepository, - request, ): agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError) machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE) @@ -229,7 +227,6 @@ def test_node_not_upserted_if_machine_retrievalerror( handler, machine_repository: IMachineRepository, node_repository: INodeRepository, - request, machine_id, ): machine_repository.get_machine_by_id = MagicMock( @@ -283,7 +280,6 @@ def test_node_not_upserted_if_machine_storageerror( handler, machine_repository: IMachineRepository, node_repository: INodeRepository, - request, ): if event == PING_SCAN_EVENT: target_machine = TARGET_MACHINE @@ -313,7 +309,6 @@ def test_failed_scan( handler, machine_repository: IMachineRepository, node_repository: INodeRepository, - request, ): machine_repository.upsert_machine = MagicMock(side_effect=StorageError) machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) From 978a2a57a97337cb06c823794115d46a1de51ddc Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:36:19 +0000 Subject: [PATCH 21/24] UT: Add default mocks for fixture methods Update machine_repository fixture to use: - machine_from_id for get_machine_by_id() - machines_from_ip for get_machines_by_ip() --- .../cc/agent_event_handlers/test_scan_event_handler.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 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 5204a30ea..a24b75ff8 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 @@ -81,6 +81,8 @@ def agent_repository() -> IAgentRepository: @pytest.fixture def machine_repository() -> IMachineRepository: machine_repository = MagicMock(spec=IMachineRepository) + machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID)) machine_repository.upsert_machine = MagicMock() return machine_repository @@ -165,7 +167,6 @@ def handler(scan_event_handler, request): indirect=["handler"], ) def test_target_machine_not_exists(event, handler, machine_repository: IMachineRepository, request): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) handler(event) @@ -187,7 +188,6 @@ def test_upserts_node( machine_repository: IMachineRepository, node_repository: INodeRepository, ): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE]) handler(event) @@ -250,9 +250,6 @@ def test_node_not_upserted_if_machine_retrievalerror( indirect=["handler"], ) def test_machine_not_upserted(event, handler, machine_repository: IMachineRepository, request): - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) - handler(event) assert not machine_repository.upsert_machine.called @@ -285,7 +282,6 @@ def test_node_not_upserted_if_machine_storageerror( target_machine = TARGET_MACHINE target_machine.operating_system = None - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) if event == PING_SCAN_EVENT: machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) @@ -311,8 +307,6 @@ def test_failed_scan( node_repository: INodeRepository, ): machine_repository.upsert_machine = MagicMock(side_effect=StorageError) - machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id) - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) handler(event) From 24684bf904ac7868c3410524a2e83f6fc451b8ec Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:46:36 +0000 Subject: [PATCH 22/24] UT: Split scan event StorageError test --- .../test_scan_event_handler.py | 35 ++++++++++--------- 1 file changed, 19 insertions(+), 16 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 a24b75ff8..8b40d0a91 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 @@ -267,27 +267,30 @@ def test_machine_not_upserted_if_existing_machine_has_os( assert not machine_repository.upsert_machine.called -@pytest.mark.parametrize( - "event,handler", - [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], - indirect=["handler"], -) -def test_node_not_upserted_if_machine_storageerror( - event, - handler, +def test_node_not_upserted_by_ping_scan_event_if_machine_storageerror( + scan_event_handler, machine_repository: IMachineRepository, node_repository: INodeRepository, ): - if event == PING_SCAN_EVENT: - target_machine = TARGET_MACHINE - target_machine.operating_system = None - - machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) - if event == PING_SCAN_EVENT: - machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) + target_machine = TARGET_MACHINE + target_machine.operating_system = None + machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip) machine_repository.upsert_machine = MagicMock(side_effect=StorageError) - handler(event) + scan_event_handler.handle_ping_scan_event(PING_SCAN_EVENT) + + assert not node_repository.upsert_communication.called + + +def test_node_not_upserted_by_tcp_scan_event_if_machine_storageerror( + scan_event_handler, + machine_repository: IMachineRepository, + node_repository: INodeRepository, +): + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + machine_repository.upsert_machine = MagicMock(side_effect=StorageError) + + scan_event_handler.handle_tcp_scan_event(TCP_SCAN_EVENT) assert not node_repository.upsert_communication.called From 28560bd65dc30d1c7141bdf76a58f6058b668774 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:49:16 +0000 Subject: [PATCH 23/24] UT: Remove more unused `request` parameters --- .../cc/agent_event_handlers/test_scan_event_handler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 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 8b40d0a91..908e17dca 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 @@ -166,7 +166,7 @@ def handler(scan_event_handler, request): [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], indirect=["handler"], ) -def test_target_machine_not_exists(event, handler, machine_repository: IMachineRepository, request): +def test_target_machine_not_exists(event, handler, machine_repository: IMachineRepository): machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) handler(event) @@ -249,14 +249,14 @@ def test_node_not_upserted_if_machine_retrievalerror( ], indirect=["handler"], ) -def test_machine_not_upserted(event, handler, machine_repository: IMachineRepository, request): +def test_machine_not_upserted(event, handler, machine_repository: IMachineRepository): handler(event) assert not machine_repository.upsert_machine.called def test_machine_not_upserted_if_existing_machine_has_os( - scan_event_handler, machine_repository: IMachineRepository, request + scan_event_handler, machine_repository: IMachineRepository ): machine_with_os = TARGET_MACHINE machine_with_os.operating_system = OperatingSystem.WINDOWS From 3b6a0cd6af84233be5b4c6b981e7d529a1bc2450 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 30 Sep 2022 19:54:52 +0000 Subject: [PATCH 24/24] UT: Split test_target_machine_not_exists --- .../test_scan_event_handler.py | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 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 908e17dca..ad1ced7fa 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 @@ -161,19 +161,28 @@ def handler(scan_event_handler, request): return getattr(scan_event_handler, request.param) -@pytest.mark.parametrize( - "event,handler", - [(PING_SCAN_EVENT, HANDLE_PING_SCAN_METHOD), (TCP_SCAN_EVENT, HANDLE_TCP_SCAN_METHOD)], - indirect=["handler"], -) -def test_target_machine_not_exists(event, handler, machine_repository: IMachineRepository): +def test_ping_scan_event_target_machine_not_exists( + scan_event_handler, machine_repository: IMachineRepository +): + event = PING_SCAN_EVENT machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) - handler(event) + scan_event_handler.handle_ping_scan_event(event) + + expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(event.target)]) + expected_machine.operating_system = event.os + machine_repository.upsert_machine.assert_called_with(expected_machine) + + +def test_tcp_scan_event_target_machine_not_exists( + scan_event_handler, machine_repository: IMachineRepository +): + event = TCP_SCAN_EVENT + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + + scan_event_handler.handle_tcp_scan_event(event) expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(event.target)]) - if event == PING_SCAN_EVENT: - expected_machine.operating_system = event.os machine_repository.upsert_machine.assert_called_with(expected_machine)