diff --git a/monkey/common/events/abstract_agent_event.py b/monkey/common/events/abstract_agent_event.py index cd8923501..760061aa1 100644 --- a/monkey/common/events/abstract_agent_event.py +++ b/monkey/common/events/abstract_agent_event.py @@ -2,11 +2,11 @@ import time from abc import ABC from ipaddress import IPv4Address from typing import FrozenSet, Union -from uuid import getnode from pydantic import Field from common.base_models import InfectionMonkeyBaseModel +from common.types import AgentID class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC): @@ -24,7 +24,7 @@ class AbstractAgentEvent(InfectionMonkeyBaseModel, ABC): :param tags: The set of tags associated with the event """ - source: int = Field(default_factory=getnode) + source: AgentID target: Union[int, IPv4Address, None] = Field(default=None) timestamp: float = Field(default_factory=time.time) tags: FrozenSet[str] = Field(default_factory=frozenset) diff --git a/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py b/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py index 10bccade3..6a9995692 100644 --- a/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py +++ b/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py @@ -6,6 +6,7 @@ from common.event_queue import IAgentEventQueue from common.events import CredentialsStolenEvent from infection_monkey.i_puppet import ICredentialCollector from infection_monkey.model import USERNAME_PREFIX +from infection_monkey.utils.ids import get_agent_id from . import pypykatz_handler from .windows_credentials import WindowsCredentials @@ -76,6 +77,7 @@ class MimikatzCredentialCollector(ICredentialCollector): def _publish_credentials_stolen_event(self, collected_credentials: Sequence[Credentials]): credentials_stolen_event = CredentialsStolenEvent( + source=get_agent_id(), tags=MIMIKATZ_EVENT_TAGS, stolen_credentials=collected_credentials, ) diff --git a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py index 8c1129455..33d2c67d8 100644 --- a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py +++ b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py @@ -11,6 +11,7 @@ from infection_monkey.telemetry.attack.t1005_telem import T1005Telem from infection_monkey.telemetry.attack.t1145_telem import T1145Telem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.utils.environment import is_windows_os +from infection_monkey.utils.ids import get_agent_id logger = logging.getLogger(__name__) @@ -172,6 +173,7 @@ def _publish_credentials_stolen_event( collected_credentials: Credentials, event_queue: IAgentEventQueue ): credentials_stolen_event = CredentialsStolenEvent( + source=get_agent_id(), tags=SSH_COLLECTOR_EVENT_TAGS, stolen_credentials=[collected_credentials], ) diff --git a/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py b/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py index a6ba4b6a7..94bf3fc14 100644 --- a/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py +++ b/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py @@ -1,13 +1,15 @@ from abc import ABC from dataclasses import dataclass +from uuid import UUID import pytest from pydantic import Field -from common.base_models import InfectionMonkeyBaseModel from common.event_serializers import IEventSerializer, PydanticEventSerializer from common.events import AbstractAgentEvent +AGENT_ID = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") + @dataclass(frozen=True) class NotAgentEvent(ABC): @@ -19,7 +21,7 @@ class SomeAgentEvent(AbstractAgentEvent): bogus: int = Field(default_factory=int) -class PydanticEvent(InfectionMonkeyBaseModel): +class PydanticEvent(AbstractAgentEvent): some_field: str @@ -29,7 +31,8 @@ def pydantic_event_serializer() -> IEventSerializer: @pytest.mark.parametrize( - "event", [NotAgentEvent(some_field=1, other_field=2.0), SomeAgentEvent(bogus=2)] + "event", + [NotAgentEvent(some_field=1, other_field=2.0), SomeAgentEvent(source=AGENT_ID, bogus=2)], ) def test_pydantic_event_serializer__serialize_wrong_type(pydantic_event_serializer, event): with pytest.raises(TypeError): @@ -42,7 +45,7 @@ def test_pydantic_event_serializer__deserialize_wrong_type(pydantic_event_serial def test_pydanitc_event_serializer__de_serialize(pydantic_event_serializer): - pydantic_event = PydanticEvent(some_field="some_field") + pydantic_event = PydanticEvent(source=AGENT_ID, some_field="some_field") serialized_event = pydantic_event_serializer.serialize(pydantic_event) deserialized_object = pydantic_event_serializer.deserialize(serialized_event)