Island, Agent: Change json encryptors to handle pydantic models
This commit is contained in:
parent
7dba3c4fed
commit
e5574240e9
|
@ -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)
|
||||||
|
|
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue