From 4bbcffabd33d09c2a39491eb0e4232545e0842b5 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 12 Jul 2022 11:44:28 +0300 Subject: [PATCH] Island: Refactor representations.py to extend JSON encoder --- .../cc/services/representations.py | 55 +++++++------------ .../cc/services/test_representations.py | 41 ++++++++------ 2 files changed, 43 insertions(+), 53 deletions(-) diff --git a/monkey/monkey_island/cc/services/representations.py b/monkey/monkey_island/cc/services/representations.py index e066764aa..fa9f2f45d 100644 --- a/monkey/monkey_island/cc/services/representations.py +++ b/monkey/monkey_island/cc/services/representations.py @@ -2,46 +2,31 @@ from datetime import datetime from enum import Enum import bson -from bson.json_util import dumps from flask import make_response +from flask.json import JSONEncoder, dumps + +from common.utils import IJSONSerializable -def _normalize_obj(obj): - if ("_id" in obj) and ("id" not in obj): - obj["id"] = obj["_id"] - del obj["_id"] - - for key, value in list(obj.items()): - obj[key] = _normalize_value(value) - return obj - - -def _normalize_value(value): - # ObjectId is serializible by default, but returns a dict - # So serialize it first into a plain string - if isinstance(value, bson.objectid.ObjectId): - return str(value) - - if isinstance(value, list): - return [_normalize_value(_value) for _value in value] - if isinstance(value, tuple): - return tuple((_normalize_value(_value) for _value in value)) - if type(value) == dict: - return _normalize_obj(value) - if isinstance(value, datetime): - return str(value) - if issubclass(type(value), Enum): - return value.name - - try: - dumps(value) - return value - except TypeError: - return value.__dict__ +class APIEncoder(JSONEncoder): + def default(self, value): + # ObjectId is serializible by default, but returns a dict + # So serialize it first into a plain string + if isinstance(value, bson.objectid.ObjectId): + return str(value) + if isinstance(value, datetime): + return str(value) + if issubclass(type(value), Enum): + return value.name + if issubclass(type(value), IJSONSerializable): + return value.__class__.to_json(value) + try: + return JSONEncoder.default(self, value) + except TypeError: + return value.__dict__ def output_json(value, code, headers=None): - value = _normalize_value(value) - resp = make_response(dumps(value), code) + resp = make_response(dumps(value, APIEncoder), code) resp.headers.extend(headers or {}) return resp diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_representations.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_representations.py index d2eef03c5..0a04032f3 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_representations.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_representations.py @@ -1,10 +1,11 @@ +import json from dataclasses import dataclass from datetime import datetime from enum import Enum import bson -from monkey_island.cc.services.representations import _normalize_value +from monkey_island.cc.services.representations import APIEncoder @dataclass @@ -17,22 +18,24 @@ bogus_object1 = MockClass(1) bogus_object2 = MockClass(2) -def test_normalize_dicts(): - assert {} == _normalize_value({}) +def test_api_encoder_dicts(): + assert json.dumps({}) == json.dumps({}, cls=APIEncoder) - assert {"a": "a"} == _normalize_value({"a": "a"}) + assert json.dumps({"a": "a"}) == json.dumps({"a": "a"}, cls=APIEncoder) - assert {"id": 12345} == _normalize_value({"id": 12345}) + assert json.dumps({"id": 12345}) == json.dumps({"id": 12345}, cls=APIEncoder) - assert {"id": obj_id_str} == _normalize_value({"id": bson.objectid.ObjectId(obj_id_str)}) + assert json.dumps({"id": obj_id_str}) == json.dumps( + {"id": bson.objectid.ObjectId(obj_id_str)}, cls=APIEncoder + ) dt = datetime.now() expected = {"a": str(dt)} - result = _normalize_value({"a": dt}) - assert expected == result + result = json.dumps({"a": dt}, cls=APIEncoder) + assert json.dumps(expected) == result -def test_normalize_complex(): +def test_api_encoder_complex(): bogus_dict = { "a": [ { @@ -44,26 +47,28 @@ def test_normalize_complex(): } expected_dict = {"a": [{"ba": obj_id_str, "bb": obj_id_str}], "b": {"id": obj_id_str}} - assert expected_dict == _normalize_value(bogus_dict) + assert json.dumps(expected_dict) == json.dumps(bogus_dict, cls=APIEncoder) -def test_normalize_list(): +def test_api_encoder_list(): bogus_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [bogus_object1]}] expected_list = [obj_id_str, {"a": "b"}, {"object": [{"a": 1}]}] - assert expected_list == _normalize_value(bogus_list) + assert json.dumps(expected_list) == json.dumps(bogus_list, cls=APIEncoder) -def test_normalize_enum(): +def test_api_encoder_enum(): class BogusEnum(Enum): bogus_val = "Bogus" my_obj = {"something": "something", "my_enum": BogusEnum.bogus_val} - assert {"something": "something", "my_enum": "bogus_val"} == _normalize_value(my_obj) + assert json.dumps({"something": "something", "my_enum": "bogus_val"}) == json.dumps( + my_obj, cls=APIEncoder + ) -def test_normalize_tuple(): - bogus_tuple = [{"my_tuple": (bogus_object1, bogus_object2, b"one_two")}] - expected_tuple = [{"my_tuple": ({"a": 1}, {"a": 2}, b"one_two")}] - assert expected_tuple == _normalize_value(bogus_tuple) +def test_api_encoder_tuple(): + bogus_tuple = [{"my_tuple": (bogus_object1, bogus_object2, "string")}] + expected_tuple = [{"my_tuple": ({"a": 1}, {"a": 2}, "string")}] + assert json.dumps(expected_tuple) == json.dumps(bogus_tuple, cls=APIEncoder)