Merge pull request #2368 from guardicore/2268-update-nodes

2268 update nodes

monkey/common/agent_events/abstract_agent_event.py

@@ -25,6 +25,6 @@ class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC):
     """
 
     source: AgentID
-    target: Union[MachineID, IPv4Address, None] = Field(default=None)
+    target: Union[IPv4Address, MachineID, None] = Field(default=None)
     timestamp: float = Field(default_factory=time.time)
     tags: FrozenSet[str] = Field(default_factory=frozenset)

monkey/common/agent_events/ping_scan_event.py

@@ -11,8 +11,9 @@ class PingScanEvent(AbstractAgentEvent):
     An event that occurs when the agent performs a ping scan on its network
 
     Attributes:
-        :param response_received: Is any response from ping recieved
-        :param os: Operating system from the target system
+        :param target: IP address of the pinged system
+        :param response_received: Indicates if target responded to the ping
+        :param os: Operating system type determined by ICMP fingerprinting
     """
 
     target: IPv4Address

monkey/common/di_container.py

@@ -1,6 +1,6 @@
 import inspect
 from contextlib import suppress
-from typing import Any, Sequence, Type, TypeVar
+from typing import Any, Sequence, Type, TypeVar, no_type_check
 
 from common.utils.code_utils import del_key
 
@@ -15,6 +15,9 @@ class UnregisteredConventionError(ValueError):
     pass
 
 
+# Mypy doesn't handle cases where abstract class is passed as Type[...]
+# https://github.com/python/mypy/issues/4717
+# We are using typing.no_type_check to mitigate these errors
 class DIContainer:
     """
     A dependency injection (DI) container that uses type annotations to resolve and inject
@@ -26,6 +29,7 @@ class DIContainer:
         self._instance_registry = {}
         self._convention_registry = {}
 
+    @no_type_check
     def register(self, interface: Type[T], concrete_type: Type[T]):
         """
         Register a concrete `type` that satisfies a given interface.
@@ -55,6 +59,7 @@ class DIContainer:
         self._type_registry[interface] = concrete_type
         del_key(self._instance_registry, interface)
 
+    @no_type_check
     def register_instance(self, interface: Type[T], instance: T):
         """
         Register a concrete instance that satisfies a given interface.
@@ -73,6 +78,7 @@ class DIContainer:
         self._instance_registry[interface] = instance
         del_key(self._type_registry, interface)
 
+    @no_type_check
     def register_convention(self, type_: Type[T], name: str, instance: T):
         """
         Register an instance as a convention
@@ -101,6 +107,7 @@ class DIContainer:
         """
         self._convention_registry[(type_, name)] = instance
 
+    @no_type_check
     def resolve(self, type_: Type[T]) -> T:
         """
         Resolves all dependencies and returns a new instance of `type_` using constructor dependency

monkey/monkey_island/cc/agent_event_handlers/__init__.py

@@ -1,2 +1,3 @@
+from .handle_ping_scan_event import handle_ping_scan_event
 from .save_event_to_event_repository import save_event_to_event_repository
 from .save_stolen_credentials_to_repository import save_stolen_credentials_to_repository

monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py

@@ -0,0 +1,76 @@
+from logging import getLogger
+
+from pydantic.json import IPv4Interface
+
+from common.agent_events import PingScanEvent
+from monkey_island.cc.models import CommunicationType, Machine
+from monkey_island.cc.repository import (
+    IAgentRepository,
+    IMachineRepository,
+    INodeRepository,
+    RetrievalError,
+    StorageError,
+    UnknownRecordError,
+)
+
+logger = getLogger(__name__)
+
+
+class handle_ping_scan_event:
+    """
+    Handles ping scan event and makes changes to Machine and Node states based on it
+    """
+
+    def __init__(
+        self,
+        agent_repository: IAgentRepository,
+        machine_repository: IMachineRepository,
+        node_repository: INodeRepository,
+    ):
+        self._agent_repository = agent_repository
+        self._machine_repository = machine_repository
+        self._node_repository = node_repository
+
+    def __call__(self, event: PingScanEvent):
+        """
+        :param event: Ping scan event to process
+        """
+        if not event.response_received:
+            return
+
+        try:
+            target_machine = self._get_target_machine(event)
+
+            self._update_target_machine_os(target_machine, event)
+            self._update_nodes(target_machine, event)
+        except (RetrievalError, StorageError, TypeError, UnknownRecordError):
+            logger.exception("Unable to process ping scan data")
+
+    def _get_target_machine(self, event: PingScanEvent) -> Machine:
+        try:
+            target_machines = self._machine_repository.get_machines_by_ip(event.target)
+            return target_machines[0]
+        except UnknownRecordError:
+            machine = Machine(
+                id=self._machine_repository.get_new_id(),
+                network_interfaces=[IPv4Interface(event.target)],
+            )
+            self._machine_repository.upsert_machine(machine)
+            return machine
+
+    def _update_target_machine_os(self, machine: Machine, event: PingScanEvent):
+        if event.os is not None and machine.operating_system is None:
+            machine.operating_system = event.os
+            self._machine_repository.upsert_machine(machine)
+
+    def _update_nodes(self, target_machine, event):
+        src_machine = self._get_source_machine(event)
+
+        # Update or create the node
+        self._node_repository.upsert_communication(
+            src_machine.id, target_machine.id, CommunicationType.SCANNED
+        )
+
+    def _get_source_machine(self, event: PingScanEvent) -> Machine:
+        agent = self._agent_repository.get_agent_by_id(event.source)
+        return self._machine_repository.get_machine_by_id(agent.machine_id)

monkey/monkey_island/cc/setup/agent_event_handlers.py

@@ -1,15 +1,23 @@
 from common import DIContainer
-from common.agent_events import CredentialsStolenEvent
+from common.agent_events import CredentialsStolenEvent, PingScanEvent
 from common.event_queue import IAgentEventQueue
 from monkey_island.cc.agent_event_handlers import (
+    handle_ping_scan_event,
     save_event_to_event_repository,
     save_stolen_credentials_to_repository,
 )
-from monkey_island.cc.repository import IAgentEventRepository, ICredentialsRepository
+from monkey_island.cc.repository import (
+    IAgentEventRepository,
+    IAgentRepository,
+    ICredentialsRepository,
+    IMachineRepository,
+    INodeRepository,
+)
 
 
 def setup_agent_event_handlers(container: DIContainer):
     _subscribe_and_store_to_event_repository(container)
+    _subscribe_ping_scan_event(container)
 
 
 def _subscribe_and_store_to_event_repository(container: DIContainer):
@@ -22,3 +30,14 @@ def _subscribe_and_store_to_event_repository(container: DIContainer):
         container.resolve(ICredentialsRepository)
     )
     agent_event_queue.subscribe_type(CredentialsStolenEvent, save_stolen_credentials_subscriber)
+
+
+def _subscribe_ping_scan_event(container: DIContainer):
+    agent_event_queue = container.resolve(IAgentEventQueue)
+    agent_repository = container.resolve(IAgentRepository)
+    machine_repository = container.resolve(IMachineRepository)
+    node_repository = container.resolve(INodeRepository)
+
+    handler = handle_ping_scan_event(agent_repository, machine_repository, node_repository)
+
+    agent_event_queue.subscribe_type(PingScanEvent, handler)

monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py

@@ -0,0 +1,268 @@
+from ipaddress import IPv4Address, IPv4Interface
+from itertools import count
+from unittest.mock import MagicMock
+from uuid import UUID
+
+import pytest
+
+from common import OperatingSystem
+from common.agent_events import PingScanEvent
+from common.types import SocketAddress
+from monkey_island.cc.agent_event_handlers import handle_ping_scan_event
+from monkey_island.cc.models import Agent, CommunicationType, Machine
+from monkey_island.cc.repository import (
+    IAgentRepository,
+    IMachineRepository,
+    INodeRepository,
+    RetrievalError,
+    StorageError,
+    UnknownRecordError,
+)
+
+SEED_ID = 99
+AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2")
+MACHINE_ID = 11
+CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000")
+AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER)
+PINGER_MACHINE = Machine(
+    id=MACHINE_ID,
+    hardware_id=5,
+    network_interfaces=[IPv4Interface("10.10.10.99/24")],
+)
+TARGET_MACHINE = Machine(
+    id=33,
+    hardware_id=9,
+    network_interfaces=[IPv4Interface("10.10.10.1/24")],
+)
+EVENT = PingScanEvent(
+    source=AGENT_ID,
+    target=IPv4Address("10.10.10.1"),
+    response_received=True,
+    os=OperatingSystem.LINUX,
+)
+
+EVENT_NO_RESPONSE = PingScanEvent(
+    source=AGENT_ID,
+    target=IPv4Address("10.10.10.1"),
+    response_received=False,
+    os=OperatingSystem.LINUX,
+)
+
+
+@pytest.fixture
+def agent_repository() -> IAgentRepository:
+    agent_repository = MagicMock(spec=IAgentRepository)
+    agent_repository.upsert_agent = MagicMock()
+    agent_repository.get_agent_by_id = MagicMock(return_value=AGENT)
+    return agent_repository
+
+
+@pytest.fixture
+def machine_repository() -> IMachineRepository:
+    machine_repository = MagicMock(spec=IMachineRepository)
+    machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID))
+    machine_repository.upsert_machine = MagicMock()
+    return machine_repository
+
+
+@pytest.fixture
+def node_repository() -> INodeRepository:
+    node_repository = MagicMock(spec=INodeRepository)
+    node_repository.upsert_communication = MagicMock()
+    return node_repository
+
+
+@pytest.fixture
+def handler(agent_repository, machine_repository, node_repository) -> handle_ping_scan_event:
+    return handle_ping_scan_event(agent_repository, machine_repository, node_repository)
+
+
+machines = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE}
+machines_by_id = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE}
+machines_by_ip = {
+    IPv4Address("10.10.10.99"): [PINGER_MACHINE],
+    IPv4Address("10.10.10.1"): [TARGET_MACHINE],
+}
+
+
+def machine_from_id(id: int):
+    return machines_by_id[id]
+
+
+def machines_from_ip(ip: IPv4Address):
+    return machines_by_ip[ip]
+
+
+class error_machine_by_id:
+    """Raise an error if the machine with the called ID matches the stored ID"""
+
+    def __init__(self, id: int, error):
+        self.id = id
+        self.error = error
+
+    def __call__(self, id: int):
+        if id == self.id:
+            raise self.error
+        else:
+            return machine_from_id(id)
+
+
+class error_machine_by_ip:
+    """Raise an error if the machine with the called IP matches the stored ID"""
+
+    def __init__(self, id: int, error):
+        self.id = id
+        self.error = error
+
+    def __call__(self, ip: IPv4Address):
+        print(f"IP is: {ip}")
+        machines = machines_from_ip(ip)
+        if machines[0].id == self.id:
+            print(f"Raise error: {self.error}")
+            raise self.error
+        else:
+            print(f"Return machine: {machines}")
+            return machines
+
+
+def test_handle_ping_scan_event__target_machine_not_exists(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+):
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError)
+
+    handler(EVENT)
+
+    expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(EVENT.target)])
+    expected_machine.operating_system = EVENT.os
+    machine_repository.upsert_machine.assert_called_with(expected_machine)
+
+
+def test_handle_ping_scan_event__target_machine_already_exists(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+):
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
+
+    handler(EVENT)
+
+    expected_machine = TARGET_MACHINE.copy()
+    expected_machine.operating_system = OperatingSystem.LINUX
+    machine_repository.upsert_machine.assert_called_with(expected_machine)
+
+
+def test_handle_ping_scan_event__upserts_node(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+):
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE])
+
+    handler(EVENT)
+
+    node_repository.upsert_communication.assert_called_with(
+        PINGER_MACHINE.id, TARGET_MACHINE.id, CommunicationType.SCANNED
+    )
+
+
+def test_handle_ping_scan_event__node_not_upserted_if_no_matching_agent(
+    handler: handle_ping_scan_event,
+    agent_repository: IAgentRepository,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+):
+    agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError)
+    machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE)
+
+    handler(EVENT)
+
+    assert not node_repository.upsert_communication.called
+
+
+def test_handle_ping_scan_event__node_not_upserted_if_no_matching_machine(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+):
+    machine_repository.get_machine_by_id = MagicMock(side_effect=UnknownRecordError)
+
+    handler(EVENT)
+
+    assert not node_repository.upsert_communication.called
+
+
+@pytest.mark.parametrize("id", [PINGER_MACHINE.id, TARGET_MACHINE.id])
+def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+    id,
+):
+    machine_repository.get_machine_by_id = MagicMock(
+        side_effect=error_machine_by_id(id, RetrievalError)
+    )
+    machine_repository.get_machines_by_ip = MagicMock(
+        side_effect=error_machine_by_ip(id, RetrievalError)
+    )
+
+    handler(EVENT)
+
+    assert not node_repository.upsert_communication.called
+
+
+def test_handle_scan_data__machine_not_upserted_if_os_is_none(
+    handler: handle_ping_scan_event, machine_repository: IMachineRepository
+):
+    event = PingScanEvent(source=EVENT.source, target=EVENT.target, response_received=True, os=None)
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
+
+    handler(event)
+
+    assert not machine_repository.upsert_machine.called
+
+
+def test_handle_scan_data__machine_not_upserted_if_existing_machine_has_os(
+    handler: handle_ping_scan_event, machine_repository: IMachineRepository
+):
+    machine_with_os = TARGET_MACHINE
+    machine_with_os.operating_system = OperatingSystem.WINDOWS
+    machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os)
+
+    handler(EVENT)
+
+    assert not machine_repository.upsert_machine.called
+
+
+def test_handle_scan_data__node_not_upserted_if_machine_storageerror(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+):
+    target_machine = TARGET_MACHINE
+    target_machine.operating_system = None
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(return_value=target_machine)
+    machine_repository.upsert_machine = MagicMock(side_effect=StorageError)
+
+    handler(EVENT)
+
+    assert not node_repository.upsert_communication.called
+
+
+def test_handle_scan_data__failed_ping(
+    handler: handle_ping_scan_event,
+    machine_repository: IMachineRepository,
+    node_repository: INodeRepository,
+):
+    machine_repository.upsert_machine = MagicMock(side_effect=StorageError)
+    machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
+    machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
+
+    handler(EVENT_NO_RESPONSE)
+
+    assert not node_repository.upsert_communication.called
+    assert not machine_repository.upsert_machine.called