diff --git a/monkey/common/events/abstract_agent_event.py b/monkey/common/events/abstract_agent_event.py index a851d9b01..cd8923501 100644 --- a/monkey/common/events/abstract_agent_event.py +++ b/monkey/common/events/abstract_agent_event.py @@ -1,13 +1,15 @@ import time from abc import ABC -from dataclasses import dataclass, field from ipaddress import IPv4Address from typing import FrozenSet, Union from uuid import getnode +from pydantic import Field -@dataclass(frozen=True) -class AbstractAgentEvent(ABC): +from common.base_models import InfectionMonkeyBaseModel + + +class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC): """ An event that was initiated or observed by an agent @@ -22,7 +24,7 @@ class AbstractAgentEvent(ABC): :param tags: The set of tags associated with the event """ - source: int = field(default_factory=getnode) - target: Union[int, IPv4Address, None] = field(default=None) - timestamp: float = field(default_factory=time.time) - tags: FrozenSet[str] = field(default_factory=frozenset) + source: int = Field(default_factory=getnode) + target: Union[int, IPv4Address, None] = Field(default=None) + timestamp: float = Field(default_factory=time.time) + tags: FrozenSet[str] = Field(default_factory=frozenset)