From da9f3c97ed271418c2854cd355a22958c3dc5b22 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 11 Jul 2022 15:38:02 +0300 Subject: [PATCH 1/3] 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) From fe6e3d75b5aa63d0a277d65d2e66342972d3235a Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 11 Jul 2022 16:53:43 +0300 Subject: [PATCH 2/3] Island, UT: Handle tuples in representations.py --- .../cc/services/representations.py | 21 +++++++++++-------- .../cc/services/test_representations.py | 19 +++++++++++------ 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/monkey/monkey_island/cc/services/representations.py b/monkey/monkey_island/cc/services/representations.py index 6267ccd8b..e066764aa 100644 --- a/monkey/monkey_island/cc/services/representations.py +++ b/monkey/monkey_island/cc/services/representations.py @@ -17,24 +17,27 @@ def _normalize_obj(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) + # 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__ - except AttributeError: - pass - - return value def output_json(value, code, headers=None): 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 6212fb892..6dfe43304 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 @@ -13,7 +13,8 @@ class MockClass: obj_id_str = "123456789012345678901234" -mock_object = MockClass(1) +mock_object1 = MockClass(1) +mock_object2 = MockClass(2) def test_normalize_dicts(): @@ -21,9 +22,9 @@ def test_normalize_dicts(): assert {"a": "a"} == _normalize_value({"a": "a"}) - assert {"id": 12345} == _normalize_value({"_id": 12345}) + assert {"id": 12345} == _normalize_value({"id": 12345}) - assert {"id": obj_id_str} == _normalize_value({"_id": bson.objectid.ObjectId(obj_id_str)}) + assert {"id": obj_id_str} == _normalize_value({"id": bson.objectid.ObjectId(obj_id_str)}) dt = datetime.now() expected = {"a": str(dt)} @@ -39,7 +40,7 @@ def test_normalize_complex(): "bb": bson.objectid.ObjectId(obj_id_str), } ], - "b": {"_id": 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}} @@ -47,16 +48,22 @@ def test_normalize_complex(): def test_normalize_list(): - mock_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [mock_object]}] + mock_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [mock_object1]}] expected_list = [obj_id_str, {"a": "b"}, {"object": [{"a": 1}]}] assert expected_list == _normalize_value(mock_list) -def test_normalize__enum(): +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) + + +def test_normalize_tuple(): + mock_tuple = [{"my_tuple": (mock_object1, mock_object2, b"one_two")}] + expected_tuple = [{"my_tuple": ({"a": 1}, {"a": 2}, b"one_two")}] + assert expected_tuple == _normalize_value(mock_tuple) From 33f15ff297eaa2a228038d9400ffd6890ee1b0c0 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 11 Jul 2022 16:55:13 +0300 Subject: [PATCH 3/3] UT: Rename mock to bogus in test_representations.py --- .../cc/services/test_representations.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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 6dfe43304..d2eef03c5 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 @@ -13,8 +13,8 @@ class MockClass: obj_id_str = "123456789012345678901234" -mock_object1 = MockClass(1) -mock_object2 = MockClass(2) +bogus_object1 = MockClass(1) +bogus_object2 = MockClass(2) def test_normalize_dicts(): @@ -33,7 +33,7 @@ def test_normalize_dicts(): def test_normalize_complex(): - mock_dict = { + bogus_dict = { "a": [ { "ba": bson.objectid.ObjectId(obj_id_str), @@ -44,14 +44,14 @@ 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(mock_dict) + assert expected_dict == _normalize_value(bogus_dict) def test_normalize_list(): - mock_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [mock_object1]}] + 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(mock_list) + assert expected_list == _normalize_value(bogus_list) def test_normalize_enum(): @@ -64,6 +64,6 @@ def test_normalize_enum(): def test_normalize_tuple(): - mock_tuple = [{"my_tuple": (mock_object1, mock_object2, b"one_two")}] + 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(mock_tuple) + assert expected_tuple == _normalize_value(bogus_tuple)