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