Merge branch 'i-node-repository' into develop

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"

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

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

monkey/monkey_island/cc/repository/i_network_map_repository.py

@@ -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

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
+        """

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

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