diff --git a/monkey/infection_monkey/telemetry/credentials_telem.py b/monkey/infection_monkey/telemetry/credentials_telem.py index 11e2bbb6d..7504b9c65 100644 --- a/monkey/infection_monkey/telemetry/credentials_telem.py +++ b/monkey/infection_monkey/telemetry/credentials_telem.py @@ -1,9 +1,8 @@ -import enum import json -from typing import Dict, Iterable +from typing import Iterable from common.common_consts.telem_categories import TelemCategoryEnum -from common.credentials import Credentials, ICredentialComponent +from common.credentials import Credentials from infection_monkey.telemetry.base_telem import BaseTelem @@ -24,24 +23,5 @@ class CredentialsTelem(BaseTelem): def send(self, log_data=True): super().send(log_data=False) - def get_data(self) -> Dict: - # TODO: At a later time we can consider factoring this into a Serializer class or similar. - return json.loads(json.dumps(self._credentials, default=_serialize)) - - -def _serialize(obj): - if isinstance(obj, enum.Enum): - return obj.name - - if isinstance(obj, ICredentialComponent): - # This is a workaround for ICredentialComponents that are implemented as dataclasses. If the - # credential_type attribute is populated with `field(init=False, ...)`, then credential_type - # is not added to the object's __dict__ attribute. The biggest risk of this workaround is - # that we might change the name of the credential_type field in ICredentialComponents, but - # automated refactoring tools would not detect that this string needs to change. This is - # mittigated by the call to getattr() below, which will raise an AttributeException if the - # attribute name changes and a unit test will fail under these conditions. - credential_type = getattr(obj, "credential_type") - return dict(obj.__dict__, **{"credential_type": credential_type}) - - return getattr(obj, "__dict__", str(obj)) + def get_data(self): + return [json.loads(Credentials.to_json(c)) for c in self._credentials] diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py index 5fca80eff..701071fbb 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_credentials_telem.py @@ -38,8 +38,7 @@ def test_credential_telem_send(spy_send_telemetry, credentials_for_test): telem = CredentialsTelem([credentials_for_test]) telem.send() - expected_data = json.dumps(expected_data, cls=telem.json_encoder) - assert spy_send_telemetry.data == expected_data + assert json.loads(spy_send_telemetry.data) == expected_data assert spy_send_telemetry.telem_category == "credentials"