Agent: Use Credentials.to_json() for CredentialsTelem serialization

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]

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"