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