Agent: Publish PingScanEvent from ping_scanner

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
 
 

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