diff --git a/monkey/monkey_island/cc/services/representations.py b/monkey/monkey_island/cc/services/representations.py index 0193fae0d..8a1e849a6 100644 --- a/monkey/monkey_island/cc/services/representations.py +++ b/monkey/monkey_island/cc/services/representations.py @@ -1,4 +1,5 @@ from datetime import datetime +from enum import Enum import bson from bson.json_util import dumps @@ -11,19 +12,27 @@ def normalize_obj(obj): del obj["_id"] for key, value in list(obj.items()): - if isinstance(value, bson.objectid.ObjectId): - obj[key] = str(value) - if isinstance(value, datetime): - obj[key] = str(value) - if isinstance(value, dict): - obj[key] = normalize_obj(value) if isinstance(value, list): for i in range(0, len(value)): - if isinstance(value[i], dict): - value[i] = normalize_obj(value[i]) + obj[key][i] = _normalize_value(value[i]) + else: + obj[key] = _normalize_value(value) return obj +def _normalize_value(value): + if type(value) == dict: + 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 + + def output_json(obj, code, headers=None): obj = normalize_obj(obj) resp = make_response(dumps(obj), code) 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 c088c3dce..e40e4470f 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,4 +1,5 @@ from datetime import datetime +from enum import Enum from unittest import TestCase import bson @@ -44,3 +45,11 @@ class TestRepresentations(TestCase): } ), ) + + def test_normalize__enum(self): + 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)