Island: Fix representations.py to encode lists and objects

Previously lists and objects didn't get encoded. This forced explicit encoding on the endpoints, which is not necessary

monkey/monkey_island/cc/services/representations.py

@@ -6,35 +6,39 @@ from bson.json_util import dumps
 from flask import make_response
 
 
-def normalize_obj(obj):
+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()):
-        if isinstance(value, list):
-            for i in range(0, len(value)):
-                obj[key][i] = _normalize_value(value[i])
-        else:
         obj[key] = _normalize_value(value)
     return obj
 
 
 def _normalize_value(value):
+    if isinstance(value, list):
+        for i in range(0, len(value)):
+            value[i] = _normalize_value(value[i])
     if type(value) == dict:
-        return normalize_obj(value)
+        return _normalize_obj(value)
     if isinstance(value, bson.objectid.ObjectId):
         return str(value)
     if isinstance(value, datetime):
         return str(value)
     if issubclass(type(value), Enum):
         return value.name
-    else:
+
+    try:
+        return value.__dict__
+    except AttributeError:
+        pass
+
     return value
 
 
-def output_json(obj, code, headers=None):
+def output_json(value, code, headers=None):
-    obj = normalize_obj(obj)
+    value = _normalize_value(value)
-    resp = make_response(dumps(obj), code)
+    resp = make_response(dumps(value), code)
     resp.headers.extend(headers or {})
     return resp

monkey/tests/unit_tests/monkey_island/cc/services/test_representations.py

@@ -1,40 +1,38 @@
+from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from unittest import TestCase
 
 import bson
 
-from monkey_island.cc.services.representations import normalize_obj
+from monkey_island.cc.services.representations import _normalize_value
 
 
-class TestRepresentations(TestCase):
+@dataclass
-    def test_normalize_obj(self):
+class MockClass:
-        # empty
+    a: int
-        self.assertEqual({}, normalize_obj({}))
 
-        # no special content
-        self.assertEqual({"a": "a"}, normalize_obj({"a": "a"}))
 
-        # _id field -> id field
-        self.assertEqual({"id": 12345}, normalize_obj({"_id": 12345}))
-
-        # obj id field -> str
 obj_id_str = "123456789012345678901234"
-        self.assertEqual(
+mock_object = MockClass(1)
-            {"id": obj_id_str}, normalize_obj({"_id": bson.objectid.ObjectId(obj_id_str)})
+
-        )
+
+def test_normalize_dicts():
+    assert {} == _normalize_value({})
+
+    assert {"a": "a"} == _normalize_value({"a": "a"})
+
+    assert {"id": 12345} == _normalize_value({"_id": 12345})
+
+    assert {"id": obj_id_str} == _normalize_value({"_id": bson.objectid.ObjectId(obj_id_str)})
 
-        # datetime -> str
     dt = datetime.now()
     expected = {"a": str(dt)}
-        result = normalize_obj({"a": dt})
+    result = _normalize_value({"a": dt})
-        self.assertEqual(expected, result)
+    assert expected == result
 
-        # dicts and lists
+
-        self.assertEqual(
+def test_normalize_complex():
-            {"a": [{"ba": obj_id_str, "bb": obj_id_str}], "b": {"id": obj_id_str}},
+    mock_dict = {
-            normalize_obj(
-                {
         "a": [
             {
                 "ba": bson.objectid.ObjectId(obj_id_str),
@@ -43,13 +41,22 @@ class TestRepresentations(TestCase):
         ],
         "b": {"_id": bson.objectid.ObjectId(obj_id_str)},
     }
-            ),
-        )
 
-    def test_normalize__enum(self):
+    expected_dict = {"a": [{"ba": obj_id_str, "bb": obj_id_str}], "b": {"id": obj_id_str}}
+    assert expected_dict == _normalize_value(mock_dict)
+
+
+def test_normalize_list():
+    mock_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [mock_object]}]
+
+    expected_list = [obj_id_str, {"a": "b"}, {"object": [{"a": 1}]}]
+    assert expected_list == _normalize_value(mock_list)
+
+
+def test_normalize__enum():
     class BogusEnum(Enum):
         bogus_val = "Bogus"
 
     my_obj = {"something": "something", "my_enum": BogusEnum.bogus_val}
 
-        assert {"something": "something", "my_enum": "bogus_val"} == normalize_obj(my_obj)
+    assert {"something": "something", "my_enum": "bogus_val"} == _normalize_value(my_obj)