diff --git a/monkey/infection_monkey/telemetry/telem_encoder.py b/monkey/infection_monkey/telemetry/telem_encoder.py index 3abe77439..224e14133 100644 --- a/monkey/infection_monkey/telemetry/telem_encoder.py +++ b/monkey/infection_monkey/telemetry/telem_encoder.py @@ -1,5 +1,7 @@ import json +from pydantic import BaseModel, SecretField + from common import OperatingSystem @@ -7,4 +9,8 @@ class TelemetryJSONEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, OperatingSystem): return obj.name + if issubclass(type(obj), BaseModel): + return obj.dict(simplify=True) + if issubclass(type(obj), SecretField): + return obj.get_secret_value() return json.JSONEncoder.default(self, obj) diff --git a/monkey/monkey_island/cc/services/representations.py b/monkey/monkey_island/cc/services/representations.py index e679231c1..19218ba37 100644 --- a/monkey/monkey_island/cc/services/representations.py +++ b/monkey/monkey_island/cc/services/representations.py @@ -5,6 +5,7 @@ from typing import Any import bson from flask import make_response +from pydantic import BaseModel from common.utils import IJSONSerializable @@ -23,6 +24,8 @@ class APIEncoder(JSONEncoder): return loads(value.__class__.to_json(value)) if issubclass(type(value), set): return list(value) + if issubclass(type(value), BaseModel): + return value.dict(simplify=True) try: return JSONEncoder.default(self, value) except TypeError: