diff --git a/monkey/common/base_models.py b/monkey/common/base_models.py index 9521e4e68..c2df68c5f 100644 --- a/monkey/common/base_models.py +++ b/monkey/common/base_models.py @@ -46,7 +46,7 @@ class InfectionMonkeyBaseModel(BaseModel): return BaseModel.dict(self, **kwargs) -class MutableBaseModel(InfectionMonkeyBaseModel): +class MutableInfectionMonkeyBaseModel(InfectionMonkeyBaseModel): class Config(InfectionMonkeyModelConfig): allow_mutation = True validate_assignment = True diff --git a/monkey/monkey_island/cc/models/agent.py b/monkey/monkey_island/cc/models/agent.py index b3e79c4e2..1740e264d 100644 --- a/monkey/monkey_island/cc/models/agent.py +++ b/monkey/monkey_island/cc/models/agent.py @@ -4,12 +4,12 @@ from uuid import UUID from pydantic import Field -from common.base_models import MutableBaseModel +from common.base_models import MutableInfectionMonkeyBaseModel from . import MachineID -class Agent(MutableBaseModel): +class Agent(MutableInfectionMonkeyBaseModel): id: UUID = Field(..., allow_mutation=False) machine_id: MachineID = Field(..., allow_mutation=False) start_time: datetime = Field(..., allow_mutation=False) diff --git a/monkey/monkey_island/cc/models/machine.py b/monkey/monkey_island/cc/models/machine.py index a7c513484..0c68f0b04 100644 --- a/monkey/monkey_island/cc/models/machine.py +++ b/monkey/monkey_island/cc/models/machine.py @@ -4,14 +4,14 @@ from typing import Optional, Sequence from pydantic import Field, PositiveInt, validator from common import OperatingSystem -from common.base_models import MutableBaseModel +from common.base_models import MutableInfectionMonkeyBaseModel from common.transforms import make_immutable_sequence from common.types import HardwareID MachineID = PositiveInt -class Machine(MutableBaseModel): +class Machine(MutableInfectionMonkeyBaseModel): id: MachineID = Field(..., allow_mutation=False) hardware_id: Optional[HardwareID] network_interfaces: Sequence[IPv4Interface] diff --git a/monkey/monkey_island/cc/models/node.py b/monkey/monkey_island/cc/models/node.py index 74eb48964..11c0784ff 100644 --- a/monkey/monkey_island/cc/models/node.py +++ b/monkey/monkey_island/cc/models/node.py @@ -2,7 +2,7 @@ from typing import Sequence, Tuple from pydantic import Field, validator -from common.base_models import MutableBaseModel +from common.base_models import MutableInfectionMonkeyBaseModel from common.transforms import make_immutable_nested_sequence from . import CommunicationType, MachineID @@ -10,7 +10,7 @@ from . import CommunicationType, MachineID ConnectionsSequence = Sequence[Tuple[MachineID, Sequence[CommunicationType]]] -class Node(MutableBaseModel): +class Node(MutableInfectionMonkeyBaseModel): machine_id: MachineID = Field(..., allow_mutation=False) connections: ConnectionsSequence