Merge pull request #2043 from guardicore/1960-deserialize-config

1960 deserialize config

monkey/common/configuration/__init__.py

@@ -1,8 +1,16 @@
-from .agent_configuration import (
-    AgentConfiguration,
-    AgentConfigurationSchema,
+from .agent_configuration import AgentConfiguration, InvalidConfigurationError
+from .agent_sub_configurations import (
+    CustomPBAConfiguration,
+    PluginConfiguration,
+    ScanTargetConfiguration,
+    ICMPScanConfiguration,
+    TCPScanConfiguration,
+    NetworkScanConfiguration,
+    ExploitationOptionsConfiguration,
+    ExploiterConfiguration,
+    ExploitationConfiguration,
+    PropagationConfiguration,
 )
 from .default_agent_configuration import (
-    DEFAULT_AGENT_CONFIGURATION_JSON,
-    build_default_agent_configuration,
+    DEFAULT_AGENT_CONFIGURATION,
 )

monkey/common/configuration/agent_configuration.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import List
+from typing import Any, List, Mapping
 
-from marshmallow import Schema, fields, post_load
+from marshmallow import Schema, fields
+from marshmallow.exceptions import MarshmallowError
 
 from .agent_sub_configuration_schemas import (
     CustomPBAConfigurationSchema,
@@ -15,6 +18,15 @@
 )
 
 
+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
@@ -24,6 +36,57 @@ class AgentConfiguration:
     payloads: List[PluginConfiguration]
     propagation: PropagationConfiguration
 
+    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_mapping(config_mapping: Mapping[str, Any]) -> AgentConfiguration:
+        """
+        Construct an AgentConfiguration from a Mapping
+
+        :param config_mapping: A Mapping that represents an AgentConfiguration
+        :return: An AgentConfiguration
+        :raises: InvalidConfigurationError if the provided Mapping does not represent a valid
+                 AgentConfiguration
+        """
+
+        try:
+            config_dict = AgentConfigurationSchema().load(config_mapping)
+            return AgentConfiguration(**config_dict)
+        except MarshmallowError as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
+
+    @staticmethod
+    def from_json(config_json: str) -> AgentConfiguration:
+        """
+        Construct an AgentConfiguration from a JSON string
+
+        :param config_json: A JSON string that represents an AgentConfiguration
+        :return: An AgentConfiguration
+        :raises: InvalidConfigurationError if the provided JSON does not represent a valid
+                 AgentConfiguration
+        """
+        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:
+        """
+        Serialize an AgentConfiguration to JSON
+
+        :param config: An AgentConfiguration
+        :return: A JSON string representing the AgentConfiguration
+        """
+        return AgentConfigurationSchema().dumps(config)
+
 
 class AgentConfigurationSchema(Schema):
     keep_tunnel_open_time = fields.Float()
@@ -32,7 +95,3 @@ class AgentConfigurationSchema(Schema):
     credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema))
     payloads = fields.List(fields.Nested(PluginConfigurationSchema))
     propagation = fields.Nested(PropagationConfigurationSchema)
-
-    @post_load
-    def _make_agent_configuration(self, data, **kwargs):
-        return AgentConfiguration(**data)