From da9f3c97ed271418c2854cd355a22958c3dc5b22 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Mon, 11 Jul 2022 15:38:02 +0300
Subject: [PATCH] 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
---
 .../cc/services/representations.py            | 28 +++---
 .../cc/services/test_representations.py       | 87 ++++++++++---------
 2 files changed, 63 insertions(+), 52 deletions(-)

diff --git a/monkey/monkey_island/cc/services/representations.py b/monkey/monkey_island/cc/services/representations.py
index 8a1e849a6..6267ccd8b 100644
--- a/monkey/monkey_island/cc/services/representations.py
+++ b/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)
+        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:
-        return value
+
+    try:
+        return value.__dict__
+    except AttributeError:
+        pass
+
+    return value
 
 
-def output_json(obj, code, headers=None):
-    obj = normalize_obj(obj)
-    resp = make_response(dumps(obj), code)
+def output_json(value, code, headers=None):
+    value = _normalize_value(value)
+    resp = make_response(dumps(value), 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 e40e4470f..6212fb892 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,55 +1,62 @@
+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):
-    def test_normalize_obj(self):
-        # empty
-        self.assertEqual({}, normalize_obj({}))
+@dataclass
+class MockClass:
+    a: int
 
-        # 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_str = "123456789012345678901234"
+mock_object = MockClass(1)
 
-        # obj id field -> str
-        obj_id_str = "123456789012345678901234"
-        self.assertEqual(
-            {"id": obj_id_str}, normalize_obj({"_id": bson.objectid.ObjectId(obj_id_str)})
-        )
 
-        # datetime -> str
-        dt = datetime.now()
-        expected = {"a": str(dt)}
-        result = normalize_obj({"a": dt})
-        self.assertEqual(expected, result)
+def test_normalize_dicts():
+    assert {} == _normalize_value({})
 
-        # dicts and lists
-        self.assertEqual(
-            {"a": [{"ba": obj_id_str, "bb": obj_id_str}], "b": {"id": obj_id_str}},
-            normalize_obj(
-                {
-                    "a": [
-                        {
-                            "ba": bson.objectid.ObjectId(obj_id_str),
-                            "bb": bson.objectid.ObjectId(obj_id_str),
-                        }
-                    ],
-                    "b": {"_id": bson.objectid.ObjectId(obj_id_str)},
-                }
-            ),
-        )
+    assert {"a": "a"} == _normalize_value({"a": "a"})
 
-    def test_normalize__enum(self):
-        class BogusEnum(Enum):
-            bogus_val = "Bogus"
+    assert {"id": 12345} == _normalize_value({"_id": 12345})
 
-        my_obj = {"something": "something", "my_enum": BogusEnum.bogus_val}
+    assert {"id": obj_id_str} == _normalize_value({"_id": bson.objectid.ObjectId(obj_id_str)})
 
-        assert {"something": "something", "my_enum": "bogus_val"} == normalize_obj(my_obj)
+    dt = datetime.now()
+    expected = {"a": str(dt)}
+    result = _normalize_value({"a": dt})
+    assert expected == result
+
+
+def test_normalize_complex():
+    mock_dict = {
+        "a": [
+            {
+                "ba": bson.objectid.ObjectId(obj_id_str),
+                "bb": bson.objectid.ObjectId(obj_id_str),
+            }
+        ],
+        "b": {"_id": bson.objectid.ObjectId(obj_id_str)},
+    }
+
+    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_value(my_obj)