diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 521e7f720..4f8c31eb1 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -11,5 +11,5 @@ from .simulation import Simulation, SimulationSchema, IslandMode from .user_credentials import UserCredentials from .machine import Machine, MachineID from .communication_type import CommunicationType -from .node import Node +from .node import Node, ConnectionTarget from .agent import Agent, AgentID diff --git a/monkey/monkey_island/cc/models/node.py b/monkey/monkey_island/cc/models/node.py index 11c0784ff..60b0349df 100644 --- a/monkey/monkey_island/cc/models/node.py +++ b/monkey/monkey_island/cc/models/node.py @@ -1,4 +1,4 @@ -from typing import Sequence, Tuple +from typing import Collection, Sequence, Tuple, TypeAlias from pydantic import Field, validator @@ -7,12 +7,12 @@ from common.transforms import make_immutable_nested_sequence from . import CommunicationType, MachineID -ConnectionsSequence = Sequence[Tuple[MachineID, Sequence[CommunicationType]]] +ConnectionTarget: TypeAlias = Tuple[MachineID, Sequence[CommunicationType]] class Node(MutableInfectionMonkeyBaseModel): machine_id: MachineID = Field(..., allow_mutation=False) - connections: ConnectionsSequence + connections: Collection[ConnectionTarget] _make_immutable_nested_sequence = validator("connections", pre=True, allow_reuse=True)( make_immutable_nested_sequence