diff --git a/monkey/common/agent_event_serializers/__init__.py b/monkey/common/agent_event_serializers/__init__.py index d42a47cb1..aa488599c 100644 --- a/monkey/common/agent_event_serializers/__init__.py +++ b/monkey/common/agent_event_serializers/__init__.py @@ -1,5 +1,5 @@ from .consts import EVENT_TYPE_FIELD -from .i_agent_event_serializer import IAgentEventSerializer, JSONSerializable +from .i_agent_event_serializer import IAgentEventSerializer from .agent_event_serializer_registry import AgentEventSerializerRegistry from .pydantic_agent_event_serializer import PydanticAgentEventSerializer from .register import register_common_agent_event_serializers diff --git a/monkey/common/agent_event_serializers/i_agent_event_serializer.py b/monkey/common/agent_event_serializers/i_agent_event_serializer.py index 348f81e04..0ec1790a7 100644 --- a/monkey/common/agent_event_serializers/i_agent_event_serializer.py +++ b/monkey/common/agent_event_serializers/i_agent_event_serializer.py @@ -1,17 +1,7 @@ from abc import ABC, abstractmethod -from typing import Dict, List, Union from common.agent_events import AbstractAgentEvent - -JSONSerializable = Union[ # type: ignore[misc] - Dict[str, "JSONSerializable"], # type: ignore[misc] - List["JSONSerializable"], # type: ignore[misc] - int, - str, - float, - bool, - None, -] +from common.types import JSONSerializable class IAgentEventSerializer(ABC): diff --git a/monkey/common/agent_event_serializers/pydantic_agent_event_serializer.py b/monkey/common/agent_event_serializers/pydantic_agent_event_serializer.py index f597cfb87..c80ae6c90 100644 --- a/monkey/common/agent_event_serializers/pydantic_agent_event_serializer.py +++ b/monkey/common/agent_event_serializers/pydantic_agent_event_serializer.py @@ -2,9 +2,10 @@ import logging from typing import Generic, Type, TypeVar from common.agent_events import AbstractAgentEvent +from common.types import JSONSerializable from common.utils.code_utils import del_key -from . import EVENT_TYPE_FIELD, IAgentEventSerializer, JSONSerializable +from . import EVENT_TYPE_FIELD, IAgentEventSerializer logger = logging.getLogger(__name__) diff --git a/monkey/common/types.py b/monkey/common/types.py index 8c5048a48..499df9339 100644 --- a/monkey/common/types.py +++ b/monkey/common/types.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass from enum import Enum from ipaddress import IPv4Address -from typing import Optional +from typing import Dict, List, Optional, Union from uuid import UUID from pydantic import ConstrainedInt, PositiveInt @@ -17,6 +17,16 @@ AgentID: TypeAlias = UUID HardwareID: TypeAlias = PositiveInt MachineID: TypeAlias = PositiveInt +JSONSerializable = Union[ # type: ignore[misc] + Dict[str, "JSONSerializable"], # type: ignore[misc] + List["JSONSerializable"], # type: ignore[misc] + int, + str, + float, + bool, + None, +] + class NetworkPort(ConstrainedInt): """ diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index eb2420c95..1fe53f25a 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -8,7 +8,7 @@ import requests from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration -from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable +from common.agent_event_serializers import AgentEventSerializerRegistry from common.agent_events import AbstractAgentEvent from common.common_consts.timeouts import ( LONG_REQUEST_TIMEOUT, @@ -16,7 +16,7 @@ from common.common_consts.timeouts import ( SHORT_REQUEST_TIMEOUT, ) from common.credentials import Credentials -from common.types import AgentID, SocketAddress +from common.types import AgentID, JSONSerializable, SocketAddress from . import ( AbstractIslandAPIClientFactory, diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py index fc347b46f..a51d47207 100644 --- a/monkey/monkey_island/cc/repository/agent_event_encryption.py +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -1,8 +1,8 @@ import json from typing import Callable -from common.agent_event_serializers import JSONSerializable from common.agent_events import AbstractAgentEvent +from common.types import JSONSerializable ENCRYPTED_PREFIX = "encrypted_" ABSTRACT_AGENT_EVENT_FIELDS = vars(AbstractAgentEvent)["__fields__"].keys()