Agent: Pass Agent ID to the scanners

Issue #3119
PR #3156

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,

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,

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,

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,

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

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

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"),
     )