Island, Agent: Change json encryptors to handle pydantic models

This commit is contained in:
vakaris_zilius 2022-09-06 11:33:38 +00:00
parent 7dba3c4fed
commit e5574240e9
2 changed files with 9 additions and 0 deletions

View File

@ -1,5 +1,7 @@
import json import json
from pydantic import BaseModel, SecretField
from common import OperatingSystem from common import OperatingSystem
@ -7,4 +9,8 @@ class TelemetryJSONEncoder(json.JSONEncoder):
def default(self, obj): def default(self, obj):
if isinstance(obj, OperatingSystem): if isinstance(obj, OperatingSystem):
return obj.name 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) return json.JSONEncoder.default(self, obj)

View File

@ -5,6 +5,7 @@ from typing import Any
import bson import bson
from flask import make_response from flask import make_response
from pydantic import BaseModel
from common.utils import IJSONSerializable from common.utils import IJSONSerializable
@ -23,6 +24,8 @@ class APIEncoder(JSONEncoder):
return loads(value.__class__.to_json(value)) return loads(value.__class__.to_json(value))
if issubclass(type(value), set): if issubclass(type(value), set):
return list(value) return list(value)
if issubclass(type(value), BaseModel):
return value.dict(simplify=True)
try: try:
return JSONEncoder.default(self, value) return JSONEncoder.default(self, value)
except TypeError: except TypeError: