Common: Switch AbstractAgentEvent.source from int to AgentID

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)

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,
         )

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],
     )

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)