From 89397d8cbd65a6285da6ad995f74f08ca09427c3 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 13:10:07 +0200 Subject: [PATCH 01/27] Agent: Rename event_queue to agent_event_queue --- monkey/infection_monkey/monkey.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index c54569360..195f82996 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -272,14 +272,14 @@ class InfectionMonkey: self._control_channel ) - event_queue = PyPubSubAgentEventQueue(Publisher()) + agent_event_queue = PyPubSubAgentEventQueue(Publisher()) self._subscribe_events( - event_queue, + agent_event_queue, propagation_credentials_repository, self._agent_event_serializer_registry, ) - puppet = self._build_puppet(event_queue) + puppet = self._build_puppet(agent_event_queue) victim_host_factory = self._build_victim_host_factory(local_network_interfaces) @@ -300,34 +300,34 @@ class InfectionMonkey: def _subscribe_events( self, - event_queue: IAgentEventQueue, + agent_event_queue: IAgentEventQueue, propagation_credentials_repository: IPropagationCredentialsRepository, agent_event_serializer_registry: AgentEventSerializerRegistry, ): - event_queue.subscribe_type( + agent_event_queue.subscribe_type( CredentialsStolenEvent, add_credentials_from_event_to_propagation_credentials_repository( propagation_credentials_repository ), ) - event_queue.subscribe_all_events( + agent_event_queue.subscribe_all_events( AgentEventForwarder(self._island_api_client, agent_event_serializer_registry).send_event ) def _build_puppet( self, - event_queue: IAgentEventQueue, + agent_event_queue: IAgentEventQueue, ) -> IPuppet: puppet = Puppet() puppet.load_plugin( "MimikatzCollector", - MimikatzCredentialCollector(event_queue), + MimikatzCredentialCollector(agent_event_queue), PluginType.CREDENTIAL_COLLECTOR, ) puppet.load_plugin( "SSHCollector", - SSHCredentialCollector(self._telemetry_messenger, event_queue), + SSHCredentialCollector(self._telemetry_messenger, agent_event_queue), PluginType.CREDENTIAL_COLLECTOR, ) @@ -341,7 +341,7 @@ class InfectionMonkey: island_api_client=self._island_api_client, ) exploit_wrapper = ExploiterWrapper( - self._telemetry_messenger, event_queue, agent_binary_repository + self._telemetry_messenger, agent_event_queue, agent_binary_repository ) puppet.load_plugin( From b1b9eb394e0e27da5aee7d784308523dfd273549 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 14:32:18 +0200 Subject: [PATCH 02/27] Common: Add default to source in AbstractAgentEvent --- monkey/common/agent_events/abstract_agent_event.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/common/agent_events/abstract_agent_event.py b/monkey/common/agent_events/abstract_agent_event.py index 8c3655602..b7719d03e 100644 --- a/monkey/common/agent_events/abstract_agent_event.py +++ b/monkey/common/agent_events/abstract_agent_event.py @@ -24,7 +24,7 @@ class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC): :param tags: The set of tags associated with the event """ - source: AgentID + source: AgentID = Field(default=None) target: Union[MachineID, IPv4Address, None] = Field(default=None) timestamp: float = Field(default_factory=time.time) tags: FrozenSet[str] = Field(default_factory=frozenset) From 799aae4498d9fe9544e15d11168f3c5dc49e8306 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 14:32:51 +0200 Subject: [PATCH 03/27] Common: Add default to PingScanEvent --- monkey/common/agent_events/ping_scan_event.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/common/agent_events/ping_scan_event.py b/monkey/common/agent_events/ping_scan_event.py index bfb650f04..b91ee9fc4 100644 --- a/monkey/common/agent_events/ping_scan_event.py +++ b/monkey/common/agent_events/ping_scan_event.py @@ -1,3 +1,5 @@ +from pydantic import Field + from common.types import PingScanData from . import AbstractAgentEvent @@ -11,4 +13,4 @@ class PingScanEvent(AbstractAgentEvent): :param scan_data: The data collected from the ping scan """ - scan_data: PingScanData + scan_data: PingScanData = Field(default=None) From 9ada95c126d90dc612e3e71650ca3d180fd80527 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 16:37:03 +0200 Subject: [PATCH 04/27] Agent: Return set when building server list --- monkey/infection_monkey/monkey.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 195f82996..cccc7af9b 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -259,9 +259,9 @@ class InfectionMonkey: return agent_event_serializer_registry def _build_server_list(self, relay_port: int): - my_servers = map(str, self._opts.servers) - relay_servers = [f"{ip}:{relay_port}" for ip in get_my_ip_addresses()] - return my_servers + relay_servers + my_servers = set(map(str, self._opts.servers)) + relay_servers = set([f"{ip}:{relay_port}" for ip in get_my_ip_addresses()]) + return my_servers.union(relay_servers) def _build_master(self, relay_port: int): servers = self._build_server_list(relay_port) From a44f763fab8040825407c44fcdf31d77d20e0276 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 20:58:54 +0200 Subject: [PATCH 05/27] Common: Remove default from AbstractAgentEvent source field --- monkey/common/agent_events/abstract_agent_event.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/common/agent_events/abstract_agent_event.py b/monkey/common/agent_events/abstract_agent_event.py index b7719d03e..8c3655602 100644 --- a/monkey/common/agent_events/abstract_agent_event.py +++ b/monkey/common/agent_events/abstract_agent_event.py @@ -24,7 +24,7 @@ class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC): :param tags: The set of tags associated with the event """ - source: AgentID = Field(default=None) + source: AgentID target: Union[MachineID, IPv4Address, None] = Field(default=None) timestamp: float = Field(default_factory=time.time) tags: FrozenSet[str] = Field(default_factory=frozenset) From acf877f3d836ea57296fa19d0eed09e08d6d6c16 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 27 Sep 2022 21:01:31 +0200 Subject: [PATCH 06/27] Common: Remove default from PingScanEvent --- monkey/common/agent_events/ping_scan_event.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/monkey/common/agent_events/ping_scan_event.py b/monkey/common/agent_events/ping_scan_event.py index b91ee9fc4..bfb650f04 100644 --- a/monkey/common/agent_events/ping_scan_event.py +++ b/monkey/common/agent_events/ping_scan_event.py @@ -1,5 +1,3 @@ -from pydantic import Field - from common.types import PingScanData from . import AbstractAgentEvent @@ -13,4 +11,4 @@ class PingScanEvent(AbstractAgentEvent): :param scan_data: The data collected from the ping scan """ - scan_data: PingScanData = Field(default=None) + scan_data: PingScanData From 5e129fd137f208331614d8e7e8c9b71aecbe4868 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 27 Sep 2022 19:19:05 -0400 Subject: [PATCH 07/27] Agent: Use set comprehension instead of set(list()) --- monkey/infection_monkey/monkey.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index cccc7af9b..f04b1b72a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -260,7 +260,7 @@ class InfectionMonkey: def _build_server_list(self, relay_port: int): my_servers = set(map(str, self._opts.servers)) - relay_servers = set([f"{ip}:{relay_port}" for ip in get_my_ip_addresses()]) + relay_servers = {f"{ip}:{relay_port}" for ip in get_my_ip_addresses()} return my_servers.union(relay_servers) def _build_master(self, relay_port: int): From 082bb3bb6f205509d4925a613356a9d9873c15ce Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 27 Sep 2022 19:20:54 -0400 Subject: [PATCH 08/27] Agent: Return a sequence from InfectionMonkey._build_server_list() --- monkey/infection_monkey/monkey.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index f04b1b72a..4b235133c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -5,7 +5,7 @@ import subprocess import sys from ipaddress import IPv4Interface from pathlib import Path, WindowsPath -from typing import List, Optional, Tuple +from typing import List, Optional, Sequence, Tuple from pubsub.core import Publisher @@ -258,10 +258,10 @@ class InfectionMonkey: return agent_event_serializer_registry - def _build_server_list(self, relay_port: int): + def _build_server_list(self, relay_port: int) -> Sequence[str]: my_servers = set(map(str, self._opts.servers)) relay_servers = {f"{ip}:{relay_port}" for ip in get_my_ip_addresses()} - return my_servers.union(relay_servers) + return list(my_servers.union(relay_servers)) def _build_master(self, relay_port: int): servers = self._build_server_list(relay_port) From ba0ffeacce3267726b37d3e8befe2be53d2cf5d6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 27 Sep 2022 19:21:34 -0400 Subject: [PATCH 09/27] Agent: Change method order --- monkey/infection_monkey/monkey.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4b235133c..98e2b46a5 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -258,11 +258,6 @@ class InfectionMonkey: return agent_event_serializer_registry - def _build_server_list(self, relay_port: int) -> Sequence[str]: - my_servers = set(map(str, self._opts.servers)) - relay_servers = {f"{ip}:{relay_port}" for ip in get_my_ip_addresses()} - return list(my_servers.union(relay_servers)) - def _build_master(self, relay_port: int): servers = self._build_server_list(relay_port) local_network_interfaces = get_network_interfaces() @@ -298,6 +293,11 @@ class InfectionMonkey: propagation_credentials_repository, ) + def _build_server_list(self, relay_port: int) -> Sequence[str]: + my_servers = set(map(str, self._opts.servers)) + relay_servers = {f"{ip}:{relay_port}" for ip in get_my_ip_addresses()} + return list(my_servers.union(relay_servers)) + def _subscribe_events( self, agent_event_queue: IAgentEventQueue, From ec56b152198ba112553c4fddae18f4b58320e067 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 11:12:06 +0200 Subject: [PATCH 10/27] Agent: Pass IAgentEventQueue to the Puppet --- monkey/infection_monkey/monkey.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 98e2b46a5..33ae633d3 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -318,7 +318,7 @@ class InfectionMonkey: self, agent_event_queue: IAgentEventQueue, ) -> IPuppet: - puppet = Puppet() + puppet = Puppet(agent_event_queue) puppet.load_plugin( "MimikatzCollector", From 0357d43d3352ee4dbf72624e60f5ab3f38b58d07 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 11:13:05 +0200 Subject: [PATCH 11/27] Agent: Accept IAgentEventQueue in Puppet constructor --- monkey/infection_monkey/puppet/puppet.py | 6 ++++-- .../unit_tests/infection_monkey/puppet/test_puppet.py | 7 ++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 80293d44d..65d4cfd0c 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -4,6 +4,7 @@ from typing import Dict, Iterable, Sequence from common.common_consts.timeouts import CONNECTION_TIMEOUT from common.credentials import Credentials +from common.event_queue import IAgentEventQueue from common.types import PingScanData from infection_monkey import network_scanning from infection_monkey.i_puppet import ( @@ -24,8 +25,9 @@ logger = logging.getLogger() class Puppet(IPuppet): - def __init__(self) -> None: + def __init__(self, agent_event_queue: IAgentEventQueue) -> None: self._plugin_registry = PluginRegistry() + self._agent_event_queue = agent_event_queue def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: self._plugin_registry.load_plugin(plugin_name, plugin, plugin_type) @@ -41,7 +43,7 @@ class Puppet(IPuppet): return pba.run(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: - return network_scanning.ping(host, timeout) + return network_scanning.ping(host, timeout, self._agent_event_queue) def scan_tcp_ports( self, host: str, ports: Sequence[int], timeout: float = CONNECTION_TIMEOUT diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py index a0df06fd6..1b0f6b9ee 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py @@ -1,13 +1,14 @@ import threading from unittest.mock import MagicMock +from common.event_queue import IAgentEventQueue from common.types import PingScanData from infection_monkey.i_puppet import PluginType from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet def test_puppet_run_payload_success(): - p = Puppet() + p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue)) payload = MagicMock() payload_name = "PayloadOne" @@ -19,7 +20,7 @@ def test_puppet_run_payload_success(): def test_puppet_run_multiple_payloads(): - p = Puppet() + p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue)) payload_1 = MagicMock() payload1_name = "PayloadOne" @@ -45,6 +46,6 @@ def test_puppet_run_multiple_payloads(): def test_fingerprint_exception_handling(monkeypatch): - p = Puppet() + p = Puppet(agent_event_queue=MagicMock(spec=IAgentEventQueue)) p._plugin_registry.get_plugin = MagicMock(side_effect=Exception) assert p.fingerprint("", "", PingScanData("windows", False), {}, {}) == EMPTY_FINGERPRINT From 228ce9bae1ef939a0edf07d0b11f16f5d5bb0186 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 11:24:25 +0200 Subject: [PATCH 12/27] Agent: Publish PingScanEvent from ping_scanner --- .../network_scanning/ping_scanner.py | 14 ++++- .../network_scanning/test_ping_scanner.py | 55 +++++++++++++------ 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index 13cafe7d1..2c2e819c0 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -4,10 +4,14 @@ import os import re import subprocess import sys +from ipaddress import IPv4Address from common import OperatingSystem +from common.agent_events import PingScanEvent +from common.event_queue import IAgentEventQueue from common.types import PingScanData from infection_monkey.utils.environment import is_windows_os +from infection_monkey.utils.ids import get_agent_id TTL_REGEX = re.compile(r"TTL=([0-9]+)\b", re.IGNORECASE) LINUX_TTL = 64 # Windows TTL is 128 @@ -17,15 +21,15 @@ EMPTY_PING_SCAN = PingScanData(False, None) logger = logging.getLogger(__name__) -def ping(host: str, timeout: float) -> PingScanData: +def ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> PingScanData: try: - return _ping(host, timeout) + return _ping(host, timeout, agent_event_queue) except Exception: logger.exception("Unhandled exception occurred while running ping") return EMPTY_PING_SCAN -def _ping(host: str, timeout: float) -> PingScanData: +def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> PingScanData: if is_windows_os(): timeout = math.floor(timeout * 1000) @@ -34,6 +38,10 @@ def _ping(host: str, timeout: float) -> PingScanData: ping_scan_data = _process_ping_command_output(ping_command_output) logger.debug(f"{host} - {ping_scan_data}") + agent_event_queue.publish( + PingScanEvent(source=get_agent_id(), target=IPv4Address(host), scan_data=ping_scan_data) + ) + return ping_scan_data diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index 9fa26456e..daa56a22b 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -5,6 +5,7 @@ from unittest.mock import MagicMock import pytest from common import OperatingSystem +from common.event_queue import IAgentEventQueue from infection_monkey.network_scanning import ping from infection_monkey.network_scanning.ping_scanner import EMPTY_PING_SCAN @@ -86,56 +87,75 @@ def set_os_windows(monkeypatch): monkeypatch.setattr("sys.platform", "win32") +@pytest.fixture +def mock_agent_event_queue(): + return MagicMock(spec=IAgentEventQueue) + + @pytest.mark.usefixtures("set_os_linux") -def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output): +def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue): patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert result.response_received assert result.os == OperatingSystem.LINUX + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.usefixtures("set_os_linux") -def test_linux_ping_no_response(patch_subprocess_running_ping_with_ping_output): +def test_linux_ping_no_response( + patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue +): patch_subprocess_running_ping_with_ping_output(LINUX_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert not result.response_received assert result.os is None + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.usefixtures("set_os_windows") -def test_windows_ping_success(patch_subprocess_running_ping_with_ping_output): +def test_windows_ping_success( + patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue +): patch_subprocess_running_ping_with_ping_output(WINDOWS_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert result.response_received assert result.os == OperatingSystem.WINDOWS + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.usefixtures("set_os_windows") -def test_windows_ping_no_response(patch_subprocess_running_ping_with_ping_output): +def test_windows_ping_no_response( + patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue +): patch_subprocess_running_ping_with_ping_output(WINDOWS_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert not result.response_received assert result.os is None + assert mock_agent_event_queue.publish.call_count == 1 -def test_malformed_ping_command_response(patch_subprocess_running_ping_with_ping_output): +def test_malformed_ping_command_response( + patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue +): patch_subprocess_running_ping_with_ping_output(MALFORMED_OUTPUT) - result = ping("192.168.1.1", 1.0) + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert not result.response_received assert result.os is None + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired") -def test_timeout_expired(): - result = ping("192.168.1.1", 1.0) +def test_timeout_expired(mock_agent_event_queue): + result = ping("192.168.1.1", 1.0, mock_agent_event_queue) assert not result.response_received assert result.os is None + assert mock_agent_event_queue.publish.call_count == 1 @pytest.fixture @@ -147,9 +167,9 @@ def ping_command_spy(monkeypatch): @pytest.fixture -def assert_expected_timeout(ping_command_spy): +def assert_expected_timeout(ping_command_spy, mock_agent_event_queue): def inner(timeout_flag, timeout_input, expected_timeout): - ping("192.168.1.1", timeout_input) + ping("192.168.1.1", timeout_input, mock_agent_event_queue) assert ping_command_spy.call_args is not None @@ -159,6 +179,8 @@ def assert_expected_timeout(ping_command_spy): timeout_flag_index = ping_command.index(timeout_flag) assert ping_command[timeout_flag_index + 1] == expected_timeout + assert mock_agent_event_queue.publish.call_count == 1 + return inner @@ -178,8 +200,9 @@ def test_linux_timeout(assert_expected_timeout): assert_expected_timeout(timeout_flag, timeout, str(math.ceil(timeout))) -def test_exception_handling(monkeypatch): +def test_exception_handling(monkeypatch, mock_agent_event_queue): monkeypatch.setattr( "infection_monkey.network_scanning.ping_scanner._ping", MagicMock(side_effect=Exception) ) - assert ping("abc", 10) == EMPTY_PING_SCAN + assert ping("abc", 10, mock_agent_event_queue) == EMPTY_PING_SCAN + assert mock_agent_event_queue.publish.call_count == 0 From 3202bfa2c1d9889da3b9858c13adb872aaaea692 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 12:12:20 +0200 Subject: [PATCH 13/27] Common: Flatten out PingScanEvent --- monkey/common/agent_events/ping_scan_event.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/monkey/common/agent_events/ping_scan_event.py b/monkey/common/agent_events/ping_scan_event.py index bfb650f04..943bf5fd2 100644 --- a/monkey/common/agent_events/ping_scan_event.py +++ b/monkey/common/agent_events/ping_scan_event.py @@ -1,4 +1,6 @@ -from common.types import PingScanData +from typing import Optional + +from common import OperatingSystem from . import AbstractAgentEvent @@ -8,7 +10,9 @@ class PingScanEvent(AbstractAgentEvent): An event that occurs when the agent performs a ping scan on its network Attributes: - :param scan_data: The data collected from the ping scan + :param response_received: Is any response from ping recieved + :param os: Operating system from the target system """ - scan_data: PingScanData + response_received: bool + os: Optional[OperatingSystem] From aa2b49bc663262e6f4175b639285e4327ee469b7 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 12:15:46 +0200 Subject: [PATCH 14/27] Agent: Seperate the generation of PingScanEvents --- .../network_scanning/ping_scanner.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index 2c2e819c0..efed79bc3 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -38,9 +38,8 @@ def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> Pin ping_scan_data = _process_ping_command_output(ping_command_output) logger.debug(f"{host} - {ping_scan_data}") - agent_event_queue.publish( - PingScanEvent(source=get_agent_id(), target=IPv4Address(host), scan_data=ping_scan_data) - ) + ping_scan_event = _generate_ping_scan_event(host, ping_scan_data) + agent_event_queue.publish(ping_scan_event) return ping_scan_data @@ -98,3 +97,12 @@ def _build_ping_command(host: str, timeout: float): # on older version of ping the timeout must be an integer, thus we use ceil return ["ping", ping_count_flag, "1", ping_timeout_flag, str(math.ceil(timeout)), host] + + +def _generate_ping_scan_event(host: str, ping_scan_data: PingScanData) -> PingScanEvent: + return PingScanEvent( + source=get_agent_id(), + target=IPv4Address(host), + response_received=ping_scan_data.response_received, + os=ping_scan_data.os, + ) From 478ea05fa991a3583e9b65f71561f69f8e34a94d Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 12:51:33 +0200 Subject: [PATCH 15/27] Common: Override target in PingScanEvent --- monkey/common/agent_events/ping_scan_event.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/common/agent_events/ping_scan_event.py b/monkey/common/agent_events/ping_scan_event.py index 943bf5fd2..0e9ee6614 100644 --- a/monkey/common/agent_events/ping_scan_event.py +++ b/monkey/common/agent_events/ping_scan_event.py @@ -1,3 +1,4 @@ +from ipaddress import IPv4Address from typing import Optional from common import OperatingSystem @@ -14,5 +15,6 @@ class PingScanEvent(AbstractAgentEvent): :param os: Operating system from the target system """ + target: IPv4Address response_received: bool os: Optional[OperatingSystem] From d235e7a19ebd777cb7e9d12aa790211fc5900d6f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:02:43 +0200 Subject: [PATCH 16/27] UT: Add unit tests for PingScanEvent --- .../agent_events/test_ping_scan_event.py | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py diff --git a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py new file mode 100644 index 000000000..173fff226 --- /dev/null +++ b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py @@ -0,0 +1,116 @@ +from ipaddress import IPv4Address +from uuid import UUID + +import pytest +from pydantic.errors import IPv4AddressError, UUIDError +from tests.unit_tests.monkey_island.cc.models.test_agent import AGENT_ID + +from common import OperatingSystem +from common.agent_events import PingScanEvent + +PING_EVENT = PingScanEvent( + source=AGENT_ID, target=IPv4Address("1.1.1.1"), response_received=True, os=OperatingSystem.LINUX +) + +PING_OBJECT_DICT = { + "source": UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10"), + "target": IPv4Address("1.1.1.1"), + "timestamp": 1664371327.4067292, + "tags": frozenset(), + "response_received": True, + "os": OperatingSystem.LINUX, +} + +PING_SIMPLE_DICT = { + "source": "012e7238-7b81-4108-8c7f-0787bc3f3c10", + "target": "1.1.1.1", + "timestamp": 1664371327.4067292, + "tags": [], + "response_received": True, + "os": "linux", +} + + +def test_constructor(): + # Raises exception_on_failure + PingScanEvent(**PING_OBJECT_DICT) + + +def test_from_dict(): + # Raises exception_on_failure + PingScanEvent(**PING_SIMPLE_DICT) + + +def test_to_dict(): + ping_scan_event = PingScanEvent(**PING_OBJECT_DICT) + + assert ping_scan_event.dict(simplify=True) == PING_SIMPLE_DICT + + +@pytest.mark.parametrize( + "key, value", + [ + ("source", "not-an-uuid"), + ("timestamp", "not-a-timestamp"), + ("response_received", "not-a-bool"), + ("os", 2.1), + ("os", "bsd"), + ], +) +def test_construct_invalid_field__type_error(key, value): + invalid_type_dict = PING_SIMPLE_DICT.copy() + invalid_type_dict[key] = value + + with pytest.raises(TypeError): + PingScanEvent(**invalid_type_dict) + + +@pytest.mark.parametrize( + "key, value, expected_error", + [ + ("source", -1, UUIDError), + ("target", "not-a-IPv4Address", IPv4AddressError), + ], +) +def test_construct_invalid_field__value_error(key, value, expected_error): + invalid_type_dict = PING_SIMPLE_DICT.copy() + invalid_type_dict[key] = value + + with pytest.raises(expected_error): + PingScanEvent(**invalid_type_dict) + + +def test_optional_os_field(): + none_field_dict = PING_SIMPLE_DICT.copy() + none_field_dict["os"] = None + + # Raises exception_on_failure + PingScanEvent(**none_field_dict) + + +def test_construct__extra_fields_forbidden(): + extra_field_dict = PING_SIMPLE_DICT.copy() + extra_field_dict["extra_field"] = 99 # red balloons + + with pytest.raises(ValueError): + PingScanEvent(**extra_field_dict) + + +def test_ping_scan_event_serialization_json(): + serialized_event = PING_EVENT.json() + assert "1.1.1.1" in serialized_event + assert "true" in serialized_event + assert OperatingSystem.LINUX.value in serialized_event + assert str(AGENT_ID) in serialized_event + + +def test_ping_scan_event_deserialization_json(): + serialized_event = PING_EVENT.json() + deserialized_event = PingScanEvent.parse_raw(serialized_event) + assert deserialized_event == PING_EVENT + + +def test_ping_scan_event_deserialization_dict(): + serialized_event = PING_EVENT.dict() + deserialized_event = PingScanEvent(**serialized_event) + assert deserialized_event == PING_EVENT From 14f80147095a56b82505035de17baf00c54527a9 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:07:38 +0200 Subject: [PATCH 17/27] Agent: Save correct event timestamp in ping_scanner --- .../infection_monkey/network_scanning/ping_scanner.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index efed79bc3..f1084f68b 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -5,6 +5,7 @@ import re import subprocess import sys from ipaddress import IPv4Address +from time import time from common import OperatingSystem from common.agent_events import PingScanEvent @@ -33,12 +34,14 @@ def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> Pin if is_windows_os(): timeout = math.floor(timeout * 1000) + event_timestamp = time() + ping_command_output = _run_ping_command(host, timeout) ping_scan_data = _process_ping_command_output(ping_command_output) logger.debug(f"{host} - {ping_scan_data}") - ping_scan_event = _generate_ping_scan_event(host, ping_scan_data) + ping_scan_event = _generate_ping_scan_event(host, ping_scan_data, event_timestamp) agent_event_queue.publish(ping_scan_event) return ping_scan_data @@ -99,10 +102,13 @@ def _build_ping_command(host: str, timeout: float): return ["ping", ping_count_flag, "1", ping_timeout_flag, str(math.ceil(timeout)), host] -def _generate_ping_scan_event(host: str, ping_scan_data: PingScanData) -> PingScanEvent: +def _generate_ping_scan_event( + host: str, ping_scan_data: PingScanData, event_timestamp: float +) -> PingScanEvent: return PingScanEvent( source=get_agent_id(), target=IPv4Address(host), + timestamp=event_timestamp, response_received=ping_scan_data.response_received, os=ping_scan_data.os, ) From d3c9691dfe7df095fcb50ed50762def629ba25fd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 15:47:58 -0400 Subject: [PATCH 18/27] UT: Add assertion to PingScanEvent serialization tests --- .../common/agent_events/test_ping_scan_event.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py index 173fff226..43f65c754 100644 --- a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py +++ b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py @@ -9,7 +9,11 @@ from common import OperatingSystem from common.agent_events import PingScanEvent PING_EVENT = PingScanEvent( - source=AGENT_ID, target=IPv4Address("1.1.1.1"), response_received=True, os=OperatingSystem.LINUX + source=AGENT_ID, + target=IPv4Address("1.1.1.1"), + timestamp=1664371327.4067292, + response_received=True, + os=OperatingSystem.LINUX, ) PING_OBJECT_DICT = { @@ -32,13 +36,11 @@ PING_SIMPLE_DICT = { def test_constructor(): - # Raises exception_on_failure - PingScanEvent(**PING_OBJECT_DICT) + assert PingScanEvent(**PING_OBJECT_DICT) == PING_EVENT def test_from_dict(): - # Raises exception_on_failure - PingScanEvent(**PING_SIMPLE_DICT) + assert PingScanEvent(**PING_SIMPLE_DICT) == PING_EVENT def test_to_dict(): From 63909938757fd9e2a102d9af80b88d1312e1b867 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 15:52:35 -0400 Subject: [PATCH 19/27] UT: Handle ValueError in test_construct_invalid_field__value_error() --- .../common/agent_events/test_ping_scan_event.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py index 43f65c754..f225b00a0 100644 --- a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py +++ b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py @@ -2,7 +2,6 @@ from ipaddress import IPv4Address from uuid import UUID import pytest -from pydantic.errors import IPv4AddressError, UUIDError from tests.unit_tests.monkey_island.cc.models.test_agent import AGENT_ID from common import OperatingSystem @@ -53,6 +52,7 @@ def test_to_dict(): "key, value", [ ("source", "not-an-uuid"), + ("source", -1), ("timestamp", "not-a-timestamp"), ("response_received", "not-a-bool"), ("os", 2.1), @@ -68,17 +68,16 @@ def test_construct_invalid_field__type_error(key, value): @pytest.mark.parametrize( - "key, value, expected_error", + "key, value", [ - ("source", -1, UUIDError), - ("target", "not-a-IPv4Address", IPv4AddressError), + ("target", "not-a-IPv4Address"), ], ) -def test_construct_invalid_field__value_error(key, value, expected_error): +def test_construct_invalid_field__value_error(key, value): invalid_type_dict = PING_SIMPLE_DICT.copy() invalid_type_dict[key] = value - with pytest.raises(expected_error): + with pytest.raises(ValueError): PingScanEvent(**invalid_type_dict) From 0cd8cd577d812d8b5b89e5ffa80f7015f1e4b63e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 15:55:07 -0400 Subject: [PATCH 20/27] UT: Remove unnecessary tests from test_ping_scan_event.py These tests are unnecessary because we will not use json() to convert to JSON. --- .../common/agent_events/test_ping_scan_event.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py index f225b00a0..c988d3e01 100644 --- a/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py +++ b/monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py @@ -97,20 +97,6 @@ def test_construct__extra_fields_forbidden(): PingScanEvent(**extra_field_dict) -def test_ping_scan_event_serialization_json(): - serialized_event = PING_EVENT.json() - assert "1.1.1.1" in serialized_event - assert "true" in serialized_event - assert OperatingSystem.LINUX.value in serialized_event - assert str(AGENT_ID) in serialized_event - - -def test_ping_scan_event_deserialization_json(): - serialized_event = PING_EVENT.json() - deserialized_event = PingScanEvent.parse_raw(serialized_event) - assert deserialized_event == PING_EVENT - - def test_ping_scan_event_deserialization_dict(): serialized_event = PING_EVENT.dict() deserialized_event = PingScanEvent(**serialized_event) From 2eee42790164655558b9ec79faab6bae1e5fbf56 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 16:02:18 -0400 Subject: [PATCH 21/27] Agent: Generate PingScanEvent timestamp closer to ping command --- .../network_scanning/ping_scanner.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index f1084f68b..c997e71e9 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -6,6 +6,7 @@ import subprocess import sys from ipaddress import IPv4Address from time import time +from typing import Tuple from common import OperatingSystem from common.agent_events import PingScanEvent @@ -34,9 +35,7 @@ def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> Pin if is_windows_os(): timeout = math.floor(timeout * 1000) - event_timestamp = time() - - ping_command_output = _run_ping_command(host, timeout) + event_timestamp, ping_command_output = _run_ping_command(host, timeout) ping_scan_data = _process_ping_command_output(ping_command_output) logger.debug(f"{host} - {ping_scan_data}") @@ -47,7 +46,7 @@ def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> Pin return ping_scan_data -def _run_ping_command(host: str, timeout: float) -> str: +def _run_ping_command(host: str, timeout: float) -> Tuple[float, str]: ping_cmd = _build_ping_command(host, timeout) logger.debug(f"Running ping command: {' '.join(ping_cmd)}") @@ -55,6 +54,8 @@ def _run_ping_command(host: str, timeout: float) -> str: # of os.device_encoding(1) will be None. Setting errors="backslashreplace" prevents a crash # in this case. See #1175 and #1403 for more information. encoding = os.device_encoding(1) + + ping_event_timestamp = time() sub_proc = subprocess.Popen( ping_cmd, stdout=subprocess.PIPE, @@ -74,9 +75,9 @@ def _run_ping_command(host: str, timeout: float) -> str: logger.debug(output) except subprocess.TimeoutExpired as te: logger.error(te) - return "" + return ping_event_timestamp, "" - return output + return ping_event_timestamp, output def _process_ping_command_output(ping_command_output: str) -> PingScanData: From a65bbc592dbaa9fc675480e02444a76639938ec7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 29 Sep 2022 15:23:57 +0530 Subject: [PATCH 22/27] UT: Check that publish is being called with expected event in test_linux_ping_success --- .../network_scanning/test_ping_scanner.py | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index daa56a22b..eb1a277ca 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -1,13 +1,17 @@ import math import subprocess from unittest.mock import MagicMock +from uuid import UUID import pytest +import infection_monkey.network_scanning.ping_scanner # noqa: F401 from common import OperatingSystem +from common.agent_events import PingScanEvent from common.event_queue import IAgentEventQueue from infection_monkey.network_scanning import ping from infection_monkey.network_scanning.ping_scanner import EMPTY_PING_SCAN +from infection_monkey.utils.ids import get_agent_id LINUX_SUCCESS_OUTPUT = """ PING 192.168.1.1 (192.168.1.1) 56(84) bytes of data. @@ -51,6 +55,19 @@ ttl=2d2! """ +@pytest.fixture(scope="module") +def patch_get_agent_id(monkeypatch): + monkeypatch.setattr("get_agent_id", lambda: UUID("9919520c-8650-4a39-aa1e-b2f4d1445159")) + + +TIMESTAMP = 123.321 + + +@pytest.fixture(autouse=True) +def patch_timestamp(monkeypatch): + monkeypatch.setattr("infection_monkey.network_scanning.ping_scanner.time", lambda: TIMESTAMP) + + @pytest.fixture def patch_subprocess_running_ping(monkeypatch): def inner(mock_obj): @@ -95,11 +112,23 @@ def mock_agent_event_queue(): @pytest.mark.usefixtures("set_os_linux") def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue): patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + host_ip = "192.168.1.1" + timeout = 1.0 + result = ping(host_ip, timeout, mock_agent_event_queue) + + event = PingScanEvent( + source=get_agent_id(), + target=host_ip, + timestamp=TIMESTAMP, + tags=frozenset(), + response_received=result.response_received, + os=result.os, + ) assert result.response_received assert result.os == OperatingSystem.LINUX assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event @pytest.mark.usefixtures("set_os_linux") From 9f15bea5bd9da311da9f85eef649b2d48997a76f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 29 Sep 2022 15:27:55 +0530 Subject: [PATCH 23/27] UT: Extract data to variables HOST_IP and TIMEOUT in test_ping_scanner.py --- .../network_scanning/test_ping_scanner.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index eb1a277ca..cac51b9bd 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -109,16 +109,17 @@ def mock_agent_event_queue(): return MagicMock(spec=IAgentEventQueue) +HOST_IP = "192.168.1.1" +TIMEOUT = 1.0 + @pytest.mark.usefixtures("set_os_linux") def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue): patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) - host_ip = "192.168.1.1" - timeout = 1.0 - result = ping(host_ip, timeout, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) event = PingScanEvent( source=get_agent_id(), - target=host_ip, + target=HOST_IP, timestamp=TIMESTAMP, tags=frozenset(), response_received=result.response_received, @@ -136,7 +137,7 @@ def test_linux_ping_no_response( patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue ): patch_subprocess_running_ping_with_ping_output(LINUX_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) assert not result.response_received assert result.os is None @@ -148,7 +149,7 @@ def test_windows_ping_success( patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue ): patch_subprocess_running_ping_with_ping_output(WINDOWS_SUCCESS_OUTPUT) - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) assert result.response_received assert result.os == OperatingSystem.WINDOWS @@ -160,7 +161,7 @@ def test_windows_ping_no_response( patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue ): patch_subprocess_running_ping_with_ping_output(WINDOWS_NO_RESPONSE_OUTPUT) - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) assert not result.response_received assert result.os is None @@ -171,7 +172,7 @@ def test_malformed_ping_command_response( patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue ): patch_subprocess_running_ping_with_ping_output(MALFORMED_OUTPUT) - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) assert not result.response_received assert result.os is None @@ -180,7 +181,7 @@ def test_malformed_ping_command_response( @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired") def test_timeout_expired(mock_agent_event_queue): - result = ping("192.168.1.1", 1.0, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) assert not result.response_received assert result.os is None @@ -198,7 +199,7 @@ def ping_command_spy(monkeypatch): @pytest.fixture def assert_expected_timeout(ping_command_spy, mock_agent_event_queue): def inner(timeout_flag, timeout_input, expected_timeout): - ping("192.168.1.1", timeout_input, mock_agent_event_queue) + ping(HOST_IP, timeout_input, mock_agent_event_queue) assert ping_command_spy.call_args is not None From d78615fa926c4017e193bb265edddbe5f8c96d2f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 29 Sep 2022 15:30:12 +0530 Subject: [PATCH 24/27] UT: Add _get_ping_scan_event()to test_ping_scanner.py --- .../network_scanning/test_ping_scanner.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index cac51b9bd..ab9c0eaa4 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -9,6 +9,7 @@ import infection_monkey.network_scanning.ping_scanner # noqa: F401 from common import OperatingSystem from common.agent_events import PingScanEvent from common.event_queue import IAgentEventQueue +from common.types import PingScanData from infection_monkey.network_scanning import ping from infection_monkey.network_scanning.ping_scanner import EMPTY_PING_SCAN from infection_monkey.utils.ids import get_agent_id @@ -112,12 +113,9 @@ def mock_agent_event_queue(): HOST_IP = "192.168.1.1" TIMEOUT = 1.0 -@pytest.mark.usefixtures("set_os_linux") -def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue): - patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) - result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) - event = PingScanEvent( +def _get_ping_scan_event(result: PingScanData): + return PingScanEvent( source=get_agent_id(), target=HOST_IP, timestamp=TIMESTAMP, @@ -126,6 +124,13 @@ def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock os=result.os, ) + +@pytest.mark.usefixtures("set_os_linux") +def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock_agent_event_queue): + patch_subprocess_running_ping_with_ping_output(LINUX_SUCCESS_OUTPUT) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) + assert result.response_received assert result.os == OperatingSystem.LINUX assert mock_agent_event_queue.publish.call_count == 1 From 03ebdd461ff90c9fc7ee4b8c4bc08e7b88587bc3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 29 Sep 2022 15:34:07 +0530 Subject: [PATCH 25/27] UT: Check that correct events are published in test_ping_scanner.py --- .../network_scanning/test_ping_scanner.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index ab9c0eaa4..c076d5b94 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -143,10 +143,12 @@ def test_linux_ping_no_response( ): patch_subprocess_running_ping_with_ping_output(LINUX_NO_RESPONSE_OUTPUT) result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event @pytest.mark.usefixtures("set_os_windows") @@ -155,10 +157,12 @@ def test_windows_ping_success( ): patch_subprocess_running_ping_with_ping_output(WINDOWS_SUCCESS_OUTPUT) result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) assert result.response_received assert result.os == OperatingSystem.WINDOWS assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event @pytest.mark.usefixtures("set_os_windows") @@ -167,10 +171,12 @@ def test_windows_ping_no_response( ): patch_subprocess_running_ping_with_ping_output(WINDOWS_NO_RESPONSE_OUTPUT) result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event def test_malformed_ping_command_response( @@ -178,19 +184,23 @@ def test_malformed_ping_command_response( ): patch_subprocess_running_ping_with_ping_output(MALFORMED_OUTPUT) result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired") def test_timeout_expired(mock_agent_event_queue): result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue) + event = _get_ping_scan_event(result) assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0] == event @pytest.fixture From 2ab86fa428a1cb1d734678a0da209acbe5f4de0c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 29 Sep 2022 07:39:57 -0400 Subject: [PATCH 26/27] UT: Remove patch_get_agent_id() 1. This fixture was never actually executing 2. get_agent_id() always returns the same value (per process). This means that _get_ping_scan_event() and the ping scanner will receive the same value for any given test. In other words, mocking it is unnecessary. --- .../infection_monkey/network_scanning/test_ping_scanner.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index c076d5b94..466fd5828 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -1,7 +1,6 @@ import math import subprocess from unittest.mock import MagicMock -from uuid import UUID import pytest @@ -56,11 +55,6 @@ ttl=2d2! """ -@pytest.fixture(scope="module") -def patch_get_agent_id(monkeypatch): - monkeypatch.setattr("get_agent_id", lambda: UUID("9919520c-8650-4a39-aa1e-b2f4d1445159")) - - TIMESTAMP = 123.321 From 66e8032ef343b3af8f33358436d1679cd4186941 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 29 Sep 2022 13:45:59 +0200 Subject: [PATCH 27/27] UT: Replace call_args comparison with assert_called_with in test_ping_scanner --- .../network_scanning/test_ping_scanner.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py index 466fd5828..102df5dd1 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py @@ -128,7 +128,7 @@ def test_linux_ping_success(patch_subprocess_running_ping_with_ping_output, mock assert result.response_received assert result.os == OperatingSystem.LINUX assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.usefixtures("set_os_linux") @@ -142,7 +142,7 @@ def test_linux_ping_no_response( assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.usefixtures("set_os_windows") @@ -156,7 +156,7 @@ def test_windows_ping_success( assert result.response_received assert result.os == OperatingSystem.WINDOWS assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.usefixtures("set_os_windows") @@ -170,7 +170,7 @@ def test_windows_ping_no_response( assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) def test_malformed_ping_command_response( @@ -183,7 +183,7 @@ def test_malformed_ping_command_response( assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired") @@ -194,7 +194,7 @@ def test_timeout_expired(mock_agent_event_queue): assert not result.response_received assert result.os is None assert mock_agent_event_queue.publish.call_count == 1 - assert mock_agent_event_queue.publish.call_args[0][0] == event + mock_agent_event_queue.publish.assert_called_with(event) @pytest.fixture