From 2ee4c9c4a5db0f0c7899f355eca2b75e4d204e2f Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 24 Mar 2023 17:08:46 +0000 Subject: [PATCH] Agent: Pass Agent ID to the scanners Issue #3119 PR #3156 --- monkey/infection_monkey/monkey.py | 4 +++- .../network_scanning/ping_scanner.py | 18 +++++++++------- .../network_scanning/tcp_scanner.py | 16 ++++++++------ monkey/infection_monkey/puppet/puppet.py | 10 ++++++--- .../network_scanning/test_ping_scanner.py | 21 ++++++++++--------- .../network_scanning/test_tcp_scanner.py | 14 ++++++------- .../infection_monkey/puppet/test_puppet.py | 2 ++ 7 files changed, 51 insertions(+), 34 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 190de0b437d..2589fd8c235 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -406,7 +406,9 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet: plugin_compatability_verifier = PluginCompatabilityVerifier( self._island_api_client, HARD_CODED_EXPLOITER_MANIFESTS ) - puppet = Puppet(self._agent_event_queue, plugin_registry, plugin_compatability_verifier) + puppet = Puppet( + self._agent_event_queue, plugin_registry, plugin_compatability_verifier, self._agent_id + ) puppet.load_plugin( AgentPluginType.CREDENTIAL_COLLECTOR, diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index 943cf0c7374..aadeebc7af9 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/monkey/infection_monkey/network_scanning/ping_scanner.py @@ -10,9 +10,9 @@ from common import OperatingSystem from common.agent_events import PingScanEvent from common.event_queue import IAgentEventQueue +from common.types import AgentID from common.utils.environment import is_windows_os from infection_monkey.i_puppet import PingScanData -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 @@ -22,15 +22,19 @@ logger = logging.getLogger(__name__) -def ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> PingScanData: +def ping( + host: str, timeout: float, agent_event_queue: IAgentEventQueue, agent_id: AgentID +) -> PingScanData: try: - return _ping(host, timeout, agent_event_queue) + return _ping(host, timeout, agent_event_queue, agent_id) except Exception: logger.exception("Unhandled exception occurred while running ping") return EMPTY_PING_SCAN -def _ping(host: str, timeout: float, agent_event_queue: IAgentEventQueue) -> PingScanData: +def _ping( + host: str, timeout: float, agent_event_queue: IAgentEventQueue, agent_id: AgentID +) -> PingScanData: if is_windows_os(): timeout = math.floor(timeout * 1000) @@ -39,7 +43,7 @@ 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}") - ping_scan_event = _generate_ping_scan_event(host, ping_scan_data, event_timestamp) + ping_scan_event = _generate_ping_scan_event(host, ping_scan_data, event_timestamp, agent_id) agent_event_queue.publish(ping_scan_event) return ping_scan_data @@ -103,10 +107,10 @@ def _build_ping_command(host: str, timeout: float): def _generate_ping_scan_event( - host: str, ping_scan_data: PingScanData, event_timestamp: float + host: str, ping_scan_data: PingScanData, event_timestamp: float, agent_id: AgentID ) -> PingScanEvent: return PingScanEvent( - source=get_agent_id(), + source=agent_id, target=IPv4Address(host), timestamp=event_timestamp, response_received=ping_scan_data.response_received, diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index 217041de798..95d28840309 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -10,10 +10,9 @@ from common.agent_events import TCPScanEvent from common.event_queue import IAgentEventQueue -from common.types import NetworkPort, PortStatus +from common.types import AgentID, NetworkPort, PortStatus from infection_monkey.i_puppet import PortScanData, PortScanDataDict from infection_monkey.network.tools import BANNER_READ, DEFAULT_TIMEOUT -from infection_monkey.utils.ids import get_agent_id logger = logging.getLogger(__name__) @@ -26,9 +25,10 @@ def scan_tcp_ports( ports_to_scan: Collection[NetworkPort], timeout: float, agent_event_queue: IAgentEventQueue, + agent_id: AgentID, ) -> PortScanDataDict: try: - return _scan_tcp_ports(host, ports_to_scan, timeout, agent_event_queue) + return _scan_tcp_ports(host, ports_to_scan, timeout, agent_event_queue, agent_id) except Exception: logger.exception("Unhandled exception occurred while trying to scan tcp ports") return EMPTY_PORT_SCAN @@ -39,24 +39,28 @@ def _scan_tcp_ports( ports_to_scan: Collection[NetworkPort], timeout: float, agent_event_queue: IAgentEventQueue, + agent_id: AgentID, ) -> PortScanDataDict: event_timestamp, open_ports = _check_tcp_ports(host, ports_to_scan, timeout) port_scan_data = _build_port_scan_data(ports_to_scan, open_ports) - tcp_scan_event = _generate_tcp_scan_event(host, port_scan_data, event_timestamp) + tcp_scan_event = _generate_tcp_scan_event(host, port_scan_data, event_timestamp, agent_id) agent_event_queue.publish(tcp_scan_event) return port_scan_data def _generate_tcp_scan_event( - host: str, port_scan_data_dict: PortScanDataDict, event_timestamp: float + host: str, + port_scan_data_dict: PortScanDataDict, + event_timestamp: float, + agent_id: AgentID, ): port_statuses = {port: psd.status for port, psd in port_scan_data_dict.items()} return TCPScanEvent( - source=get_agent_id(), + source=agent_id, target=IPv4Address(host), timestamp=event_timestamp, ports=port_statuses, diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 11313f859e5..32a39633a8d 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -5,7 +5,7 @@ from common.common_consts.timeouts import CONNECTION_TIMEOUT from common.credentials import Credentials from common.event_queue import IAgentEventQueue -from common.types import Event, NetworkPort +from common.types import AgentID, Event, NetworkPort from infection_monkey import network_scanning from infection_monkey.i_puppet import ( ExploiterResultData, @@ -31,10 +31,12 @@ def __init__( agent_event_queue: IAgentEventQueue, plugin_registry: PluginRegistry, plugin_compatability_verifier: PluginCompatabilityVerifier, + agent_id: AgentID, ) -> None: self._plugin_registry = plugin_registry self._agent_event_queue = agent_event_queue self._plugin_compatability_verifier = plugin_compatability_verifier + self._agent_id = agent_id def load_plugin(self, plugin_type: AgentPluginType, plugin_name: str, plugin: object) -> None: self._plugin_registry.load_plugin(plugin_type, plugin_name, plugin) @@ -46,12 +48,14 @@ def run_credential_collector(self, name: str, options: Dict) -> Sequence[Credent return credential_collector.collect_credentials(options) def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData: - return network_scanning.ping(host, timeout, self._agent_event_queue) + return network_scanning.ping(host, timeout, self._agent_event_queue, self._agent_id) def scan_tcp_ports( self, host: str, ports: Sequence[NetworkPort], timeout: float = CONNECTION_TIMEOUT ) -> PortScanDataDict: - return network_scanning.scan_tcp_ports(host, ports, timeout, self._agent_event_queue) + return network_scanning.scan_tcp_ports( + host, ports, timeout, self._agent_event_queue, self._agent_id + ) def fingerprint( self, 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 850c514457f..925202f8610 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 @@ -7,10 +7,10 @@ import infection_monkey.network_scanning.ping_scanner # noqa: F401 from common import OperatingSystem from common.agent_events import PingScanEvent +from common.types import AgentID from infection_monkey.i_puppet 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. @@ -55,6 +55,7 @@ TIMESTAMP = 123.321 +AGENT_ID = AgentID("a15111d3-e150-4ad5-a9b6-34f9d3a3105b") @pytest.fixture(autouse=True) @@ -104,7 +105,7 @@ def set_os_windows(monkeypatch): def _get_ping_scan_event(result: PingScanData): return PingScanEvent( - source=get_agent_id(), + source=AGENT_ID, target=HOST_IP, timestamp=TIMESTAMP, tags=frozenset(), @@ -116,7 +117,7 @@ def _get_ping_scan_event(result: PingScanData): @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) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert result.response_received @@ -130,7 +131,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(HOST_IP, TIMEOUT, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert not result.response_received @@ -144,7 +145,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(HOST_IP, TIMEOUT, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert result.response_received @@ -158,7 +159,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(HOST_IP, TIMEOUT, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert not result.response_received @@ -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(HOST_IP, TIMEOUT, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert not result.response_received @@ -182,7 +183,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(HOST_IP, TIMEOUT, mock_agent_event_queue) + result = ping(HOST_IP, TIMEOUT, mock_agent_event_queue, AGENT_ID) event = _get_ping_scan_event(result) assert not result.response_received @@ -202,7 +203,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(HOST_IP, timeout_input, mock_agent_event_queue) + ping(HOST_IP, timeout_input, mock_agent_event_queue, AGENT_ID) assert ping_command_spy.call_args is not None @@ -237,5 +238,5 @@ 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, mock_agent_event_queue) == EMPTY_PING_SCAN + assert ping("abc", 10, mock_agent_event_queue, AGENT_ID) == EMPTY_PING_SCAN assert mock_agent_event_queue.publish.call_count == 0 diff --git a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py index 2fd67590d46..38621cd0ca1 100644 --- a/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py +++ b/monkey/tests/unit_tests/infection_monkey/network_scanning/test_tcp_scanner.py @@ -3,11 +3,10 @@ import pytest from common.agent_events import TCPScanEvent -from common.types import PortStatus +from common.types import AgentID, PortStatus from infection_monkey.i_puppet import PortScanData from infection_monkey.network_scanning import scan_tcp_ports from infection_monkey.network_scanning.tcp_scanner import EMPTY_PORT_SCAN -from infection_monkey.utils.ids import get_agent_id PORTS_TO_SCAN = [22, 80, 8080, 143, 445, 2222] @@ -16,6 +15,7 @@ TIMESTAMP = 123.321 HOST_IP = "127.0.0.1" +AGENT_ID = AgentID("b63e5ca3-e33b-4c3b-96d3-2e6f10d6e2d9") @pytest.fixture(autouse=True) @@ -38,7 +38,7 @@ def _get_tcp_scan_event(port_scan_data: PortScanData): port_statuses = {port: psd.status for port, psd in port_scan_data.items()} return TCPScanEvent( - source=get_agent_id(), + source=AGENT_ID, target=HOST_IP, timestamp=TIMESTAMP, ports=port_statuses, @@ -51,7 +51,7 @@ def test_tcp_successful( ): closed_ports = [8080, 143, 445] - port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue, AGENT_ID) assert len(port_scan_data) == 6 for port in open_ports_data.keys(): @@ -74,7 +74,7 @@ def test_tcp_successful( def test_tcp_empty_response( monkeypatch, patch_check_tcp_ports, open_ports_data, mock_agent_event_queue ): - port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, PORTS_TO_SCAN, 0, mock_agent_event_queue, AGENT_ID) assert len(port_scan_data) == 6 for port in open_ports_data: @@ -92,7 +92,7 @@ def test_tcp_empty_response( def test_tcp_no_ports_to_scan( monkeypatch, patch_check_tcp_ports, open_ports_data, mock_agent_event_queue ): - port_scan_data = scan_tcp_ports(HOST_IP, [], 0, mock_agent_event_queue) + port_scan_data = scan_tcp_ports(HOST_IP, [], 0, mock_agent_event_queue, AGENT_ID) assert len(port_scan_data) == 0 @@ -107,5 +107,5 @@ def test_exception_handling(monkeypatch, mock_agent_event_queue): "infection_monkey.network_scanning.tcp_scanner._scan_tcp_ports", MagicMock(side_effect=Exception), ) - assert scan_tcp_ports("abc", [123], 123, mock_agent_event_queue) == EMPTY_PORT_SCAN + assert scan_tcp_ports("abc", [123], 123, mock_agent_event_queue, AGENT_ID) == EMPTY_PORT_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 1b575138c3d..41043d121e0 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py @@ -8,6 +8,7 @@ from common import OperatingSystem from common.agent_plugins import AgentPluginType from common.event_queue import IAgentEventQueue +from common.types import AgentID from infection_monkey.i_puppet import IncompatibleOperatingSystemError, PingScanData, TargetHost from infection_monkey.puppet import PluginCompatabilityVerifier, PluginRegistry from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet @@ -51,6 +52,7 @@ def puppet( agent_event_queue=mock_agent_event_queue, plugin_registry=mock_plugin_registry, plugin_compatability_verifier=mock_plugin_compatability_verifier, + agent_id=AgentID("4277aa81-660b-4673-b96c-443ed525b4d0"), )