diff --git a/monkey/monkey_island/cc/models/communication_type.py b/monkey/monkey_island/cc/models/communication_type.py index 0991dd5de..d7b090143 100644 --- a/monkey/monkey_island/cc/models/communication_type.py +++ b/monkey/monkey_island/cc/models/communication_type.py @@ -12,4 +12,3 @@ class CommunicationType(Enum): SCANNED = "scanned" EXPLOITED = "exploited" CC = "cc" - CC_TUNNEL = "cc_tunnel" diff --git a/monkey/monkey_island/cc/models/node.py b/monkey/monkey_island/cc/models/node.py index 11c0784ff..954869c5f 100644 --- a/monkey/monkey_island/cc/models/node.py +++ b/monkey/monkey_island/cc/models/node.py @@ -1,19 +1,25 @@ -from typing import Sequence, Tuple +from typing import Mapping, Tuple, TypeAlias from pydantic import Field, validator from common.base_models import MutableInfectionMonkeyBaseModel -from common.transforms import make_immutable_nested_sequence from . import CommunicationType, MachineID -ConnectionsSequence = Sequence[Tuple[MachineID, Sequence[CommunicationType]]] +NodeConnections: TypeAlias = Mapping[MachineID, Tuple[CommunicationType, ...]] class Node(MutableInfectionMonkeyBaseModel): - machine_id: MachineID = Field(..., allow_mutation=False) - connections: ConnectionsSequence + """ + A network node and its outbound connections/communications - _make_immutable_nested_sequence = validator("connections", pre=True, allow_reuse=True)( - make_immutable_nested_sequence - ) + A node is identified by a MachineID and tracks all outbound communication to other machines on + the network. This is particularly useful for creating graphs of Infection Monkey's activity + throughout the network. + """ + + machine_id: MachineID = Field(..., allow_mutation=False) + """The MachineID of the node (source)""" + + connections: NodeConnections + """All outbound connections from this node to other machines""" diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index e1d5fd47b..6bda93bbf 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -9,6 +9,7 @@ from .i_credentials_repository import ICredentialsRepository from .i_user_repository import IUserRepository from .i_machine_repository import IMachineRepository from .i_agent_repository import IAgentRepository +from .i_node_repository import INodeRepository from .local_storage_file_repository import LocalStorageFileRepository diff --git a/monkey/monkey_island/cc/repository/i_network_map_repository.py b/monkey/monkey_island/cc/repository/i_network_map_repository.py deleted file mode 100644 index 0b0b29847..000000000 --- a/monkey/monkey_island/cc/repository/i_network_map_repository.py +++ /dev/null @@ -1,11 +0,0 @@ -from abc import ABC - - -class INetworkMapRepository(ABC): - - # TODO Define NetMap object - def get_map(self) -> NetMap: # noqa: F821 - pass - - def save_netmap(self, netmap: NetMap): # noqa: F821 - pass diff --git a/monkey/monkey_island/cc/repository/i_node_repository.py b/monkey/monkey_island/cc/repository/i_node_repository.py new file mode 100644 index 000000000..b8ea0c49d --- /dev/null +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -0,0 +1,35 @@ +from abc import ABC, abstractmethod +from typing import Sequence + +from monkey_island.cc.models import CommunicationType, MachineID, Node + + +class INodeRepository(ABC): + """A repository used to store and retrieve `Node` objects""" + + @abstractmethod + def upsert_communication( + self, src: MachineID, dst: MachineID, communication_type: CommunicationType + ): + """ + Insert or update a node connection + + Insert or update data about how network nodes are able to communicate. Nodes are identified + by MachineID and store information about all outbound connections to other machines. By + providing a source machine, target machine, and how they communicated, nodes in this + repository can be created if they don't exist or updated if they do. + + :param src: The machine that the connection or communication originated from + :param dst: The machine that the src communicated with + :param communication_type: The way the machines communicated + :raises StorageError: If an error occurred while attempting to upsert the Node + """ + + @abstractmethod + def get_nodes(self) -> Sequence[Node]: + """ + Return all nodes that are stored in the repository + + :return: All known Nodes + :raises RetrievalError: If an error occurred while attempting to retrieve the nodes + """ diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_node.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_node.py index e980ad5d6..fdb0e6862 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_node.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_node.py @@ -1,16 +1,16 @@ -from typing import MutableSequence +from typing import MutableMapping, MutableSequence import pytest -from monkey_island.cc.models import CommunicationType, Node +from monkey_island.cc.models import CommunicationType, MachineID, Node def test_constructor(): machine_id = 1 - connections = ( - (6, (CommunicationType.SCANNED,)), - (7, (CommunicationType.SCANNED, CommunicationType.EXPLOITED)), - ) + connections = { + 6: (CommunicationType.SCANNED,), + 7: (CommunicationType.SCANNED, CommunicationType.EXPLOITED), + } n = Node( machine_id=1, connections=connections, @@ -23,13 +23,10 @@ def test_constructor(): def test_serialization(): node_dict = { "machine_id": 1, - "connections": [ - [ - 6, - ["cc", "scanned"], - ], - [7, ["exploited", "cc_tunnel"]], - ], + "connections": { + "6": ["cc", "scanned"], + "7": ["exploited", "cc"], + }, } n = Node(**node_dict) @@ -37,7 +34,7 @@ def test_serialization(): def test_machine_id_immutable(): - n = Node(machine_id=1, connections=[]) + n = Node(machine_id=1, connections={}) with pytest.raises(TypeError): n.machine_id = 2 @@ -45,52 +42,45 @@ def test_machine_id_immutable(): def test_machine_id__invalid_type(): with pytest.raises(TypeError): - Node(machine_id=None, connections=[]) + Node(machine_id=None, connections={}) def test_machine_id__invalid_value(): with pytest.raises(ValueError): - Node(machine_id=-5, connections=[]) + Node(machine_id=-5, connections={}) def test_connections__mutable(): - n = Node(machine_id=1, connections=[]) + n = Node(machine_id=1, connections={}) # Raises exception on failure - n.connections = [(5, []), (7, [])] + n.connections = {5: [], 7: []} def test_connections__invalid_machine_id(): - n = Node(machine_id=1, connections=[]) + n = Node(machine_id=1, connections={}) with pytest.raises(ValueError): - n.connections = [(5, []), (-5, [])] + n.connections = {5: [], -5: []} def test_connections__recursively_immutable(): n = Node( machine_id=1, - connections=[ - [6, [CommunicationType.SCANNED]], - [7, [CommunicationType.SCANNED, CommunicationType.EXPLOITED]], - ], + connections={ + 6: [CommunicationType.SCANNED], + 7: [CommunicationType.SCANNED, CommunicationType.EXPLOITED], + }, ) - assert not isinstance(n.connections, MutableSequence) - assert not isinstance(n.connections[0], MutableSequence) - assert not isinstance(n.connections[1], MutableSequence) - assert not isinstance(n.connections[0][1], MutableSequence) - assert not isinstance(n.connections[1][1], MutableSequence) + for connections in n.connections.values(): + assert not isinstance(connections, MutableSequence) def test_connections__set_invalid_communications_type(): - connections = ( - [ - [8, [CommunicationType.SCANNED, "invalid_comm_type"]], - ], - ) + connections = {8: [CommunicationType.SCANNED, "invalid_comm_type"]} - n = Node(machine_id=1, connections=[]) + n = Node(machine_id=1, connections={}) with pytest.raises(ValueError): n.connections = connections diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 14a311f0e..a8369b7da 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -18,7 +18,6 @@ from monkey_island.cc.repository.i_attack_repository import IAttackRepository from monkey_island.cc.repository.i_config_repository import IConfigRepository from monkey_island.cc.repository.i_log_repository import ILogRepository from monkey_island.cc.repository.i_machine_repository import IMachineRepository -from monkey_island.cc.repository.i_network_map_repository import INetworkMapRepository from monkey_island.cc.repository.i_report_repository import IReportRepository from monkey_island.cc.repository.i_simulation_repository import ISimulationRepository from monkey_island.cc.repository.i_telemetry_repository import ITelemetryRepository @@ -260,8 +259,10 @@ IMachineRepository.upsert_machine IMachineRepository.get_machine_by_id IMachineRepository.get_machine_by_hardware_id IMachineRepository.get_machines_by_ip -INetworkMapRepository.get_map -INetworkMapRepository.save_netmap +INodeRepository +INodeRepository.upsert_communication +INodeRepository.communication_type +INodeRepository.get_nodes IReportRepository ISimulationRepository.save_simulation ISimulationRepository.get_simulation