From e5574240e9f3c2c93497ce18384aa64382a181ed Mon Sep 17 00:00:00 2001
From: vakaris_zilius <vakarisz@yahoo.com>
Date: Tue, 6 Sep 2022 11:33:38 +0000
Subject: [PATCH] Island, Agent: Change json encryptors to handle pydantic
 models

---
 monkey/infection_monkey/telemetry/telem_encoder.py  | 6 ++++++
 monkey/monkey_island/cc/services/representations.py | 3 +++
 2 files changed, 9 insertions(+)

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: