diff --git a/monkey/common/agent_events/ping_scan_event.py b/monkey/common/agent_events/ping_scan_event.py index bfb650f04..0e9ee6614 100644 --- a/monkey/common/agent_events/ping_scan_event.py +++ b/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] diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index c54569360..33ae633d3 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,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( diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index 13cafe7d1..c997e71e9 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/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, + ) 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/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..c988d3e01 --- /dev/null +++ b/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 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..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 @@ -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(): - result = ping("192.168.1.1", 1.0) +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 + 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 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