Merge pull request #2358 from guardicore/2268-modify-ping-scanner-with-agent-queue

2268 modify ping scanner with agent queue

monkey/common/agent_events/ping_scan_event.py

@@ -1,4 +1,7 @@
-from common.types import PingScanData
+from ipaddress import IPv4Address
+from typing import Optional
+
+from common import OperatingSystem
 
 from . import AbstractAgentEvent
 
@@ -8,7 +11,10 @@ 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
+    target: IPv4Address
+    response_received: bool
+    os: Optional[OperatingSystem]

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,11 +258,6 @@ 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
-
     def _build_master(self, relay_port: int):
         servers = self._build_server_list(relay_port)
         local_network_interfaces = get_network_interfaces()
@@ -272,14 +267,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)
 
@@ -298,36 +293,41 @@ 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,
-        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 = Puppet(agent_event_queue)
 
         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(

monkey/infection_monkey/network_scanning/ping_scanner.py

@@ -4,10 +4,16 @@ import os
 import re
 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
+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,27 +23,30 @@ 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)
 
-    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}")
 
+    ping_scan_event = _generate_ping_scan_event(host, ping_scan_data, event_timestamp)
+    agent_event_queue.publish(ping_scan_event)
+
     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)}")
 
@@ -45,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,
@@ -64,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:
@@ -90,3 +101,15 @@ 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, 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,
+    )

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

monkey/tests/unit_tests/common/agent_events/test_ping_scan_event.py

@@ -0,0 +1,103 @@
+from ipaddress import IPv4Address
+from uuid import UUID
+
+import pytest
+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"),
+    timestamp=1664371327.4067292,
+    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():
+    assert PingScanEvent(**PING_OBJECT_DICT) == PING_EVENT
+
+
+def test_from_dict():
+    assert PingScanEvent(**PING_SIMPLE_DICT) == PING_EVENT
+
+
+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"),
+        ("source", -1),
+        ("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",
+    [
+        ("target", "not-a-IPv4Address"),
+    ],
+)
+def test_construct_invalid_field__value_error(key, value):
+    invalid_type_dict = PING_SIMPLE_DICT.copy()
+    invalid_type_dict[key] = value
+
+    with pytest.raises(ValueError):
+        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_deserialization_dict():
+    serialized_event = PING_EVENT.dict()
+    deserialized_event = PingScanEvent(**serialized_event)
+    assert deserialized_event == PING_EVENT

monkey/tests/unit_tests/infection_monkey/network_scanning/test_ping_scanner.py

@@ -4,9 +4,14 @@ from unittest.mock import MagicMock
 
 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 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
 
 LINUX_SUCCESS_OUTPUT = """
 PING 192.168.1.1 (192.168.1.1) 56(84) bytes of data.
@@ -50,6 +55,14 @@ ttl=2d2!
 """
 
 
+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):
@@ -86,56 +99,102 @@ def set_os_windows(monkeypatch):
     monkeypatch.setattr("sys.platform", "win32")
 
 
+@pytest.fixture
+def mock_agent_event_queue():
+    return MagicMock(spec=IAgentEventQueue)
+
+
+HOST_IP = "192.168.1.1"
+TIMEOUT = 1.0
+
+
+def _get_ping_scan_event(result: PingScanData):
+    return PingScanEvent(
+        source=get_agent_id(),
+        target=HOST_IP,
+        timestamp=TIMESTAMP,
+        tags=frozenset(),
+        response_received=result.response_received,
+        os=result.os,
+    )
+
+
 @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(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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
 @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(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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
 @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(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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
 @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(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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
-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(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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
 @pytest.mark.usefixtures("patch_subprocess_running_ping_to_raise_timeout_expired")
-def test_timeout_expired():
+def test_timeout_expired(mock_agent_event_queue):
-    result = ping("192.168.1.1", 1.0)
+    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
+    mock_agent_event_queue.publish.assert_called_with(event)
 
 
 @pytest.fixture
@@ -147,9 +206,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(HOST_IP, timeout_input, mock_agent_event_queue)
 
         assert ping_command_spy.call_args is not None
 
@@ -159,6 +218,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 +239,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

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