Common: Encapsulate MarshmallowError

monkey/common/configuration/agent_configuration.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from typing import List
 
 from marshmallow import Schema, fields
+from marshmallow.exceptions import MarshmallowError
 
 from .agent_sub_configuration_schemas import (
     CustomPBAConfigurationSchema,
@@ -21,6 +22,11 @@ class InvalidConfigurationError(Exception):
     pass
 
 
+INVALID_CONFIGURATION_ERROR_MESSAGE = (
+    "Cannot construct an AgentConfiguration object with the supplied, invalid data:"
+)
+
+
 @dataclass(frozen=True)
 class AgentConfiguration:
     keep_tunnel_open_time: float
@@ -33,17 +39,26 @@ class AgentConfiguration:
     def __post_init__(self):
         # This will raise an exception if the object is invalid. Calling this in __post__init()
         # makes it impossible to construct an invalid object
+        try:
             AgentConfigurationSchema().dump(self)
+        except Exception as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
 
     @staticmethod
     def from_dict(dict_: dict):
+        try:
             config_dict = AgentConfigurationSchema().load(dict_)
             return AgentConfiguration(**config_dict)
+        except MarshmallowError as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
 
     @staticmethod
     def from_json(config_json: dict):
+        try:
             config_dict = AgentConfigurationSchema().loads(config_json)
             return AgentConfiguration(**config_dict)
+        except MarshmallowError as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
 
     @staticmethod
     def to_json(config: AgentConfiguration) -> str:

monkey/monkey_island/cc/resources/agent_configuration.py

@@ -1,9 +1,9 @@
 import json
 
-import marshmallow
 from flask import make_response, request
 
 from common.configuration.agent_configuration import AgentConfiguration as AgentConfigurationObject
+from common.configuration.agent_configuration import InvalidConfigurationError
 from monkey_island.cc.repository import IAgentConfigurationRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
@@ -28,7 +28,7 @@ class AgentConfiguration(AbstractResource):
             configuration_object = AgentConfigurationObject.from_json(request.data)
             self._agent_configuration_repository.store_configuration(configuration_object)
             return make_response({}, 200)
-        except (marshmallow.exceptions.ValidationError, json.JSONDecodeError) as err:
+        except (InvalidConfigurationError, json.JSONDecodeError) as err:
             return make_response(
                 {"message": f"Invalid configuration supplied: {err}"},
                 400,

monkey/tests/unit_tests/common/configuration/test_agent_configuration.py

@@ -26,7 +26,11 @@ from tests.common.example_agent_configuration import (
     WINDOWS_FILENAME,
 )
 
-from common.configuration import DEFAULT_AGENT_CONFIGURATION_JSON, AgentConfiguration
+from common.configuration import (
+    DEFAULT_AGENT_CONFIGURATION_JSON,
+    AgentConfiguration,
+    InvalidConfigurationError,
+)
 from common.configuration.agent_configuration import AgentConfigurationSchema
 from common.configuration.agent_sub_configuration_schemas import (
     CustomPBAConfigurationSchema,
@@ -172,7 +176,7 @@ def test_agent_configuration():
 
 def test_incorrect_type():
     valid_config = AgentConfiguration.from_dict(AGENT_CONFIGURATION)
-    with pytest.raises(Exception):
+    with pytest.raises(InvalidConfigurationError):
         valid_config_dict = valid_config.__dict__
         valid_config_dict["keep_tunnel_open_time"] = "not_a_float"
         AgentConfiguration(**valid_config_dict)
@@ -193,6 +197,14 @@ def test_from_dict():
     assert schema.dump(config) == dict_
 
 
+def test_from_dict__invalid_data():
+    dict_ = json.loads(DEFAULT_AGENT_CONFIGURATION_JSON)
+    dict_["payloads"] = "payloads"
+
+    with pytest.raises(InvalidConfigurationError):
+        AgentConfiguration.from_dict(dict_)
+
+
 def test_from_json():
     schema = AgentConfigurationSchema()
     dict_ = json.loads(DEFAULT_AGENT_CONFIGURATION_JSON)
@@ -202,6 +214,14 @@ def test_from_json():
     assert schema.dump(config) == dict_
 
 
+def test_from_json__invalid_data():
+    invalid_dict = json.loads(DEFAULT_AGENT_CONFIGURATION_JSON)
+    invalid_dict["payloads"] = "payloads"
+
+    with pytest.raises(InvalidConfigurationError):
+        AgentConfiguration.from_json(json.dumps(invalid_dict))
+
+
 def test_to_json():
     config = AgentConfiguration.from_json(DEFAULT_AGENT_CONFIGURATION_JSON)