Island: Refactor representations.py to extend JSON encoder

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

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)