Merge pull request #2131 from guardicore/2004-agent-configuration-validation

AgentConfiguration docstring + validation

monkey/common/agent_configuration/agent_configuration.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Mapping, Tuple
 
-from marshmallow import Schema, fields
+from marshmallow import Schema, fields, validate
 from marshmallow.exceptions import MarshmallowError
 
 from ..utils.code_utils import freeze_lists_in_mapping
@@ -32,6 +32,19 @@ class InvalidConfigurationError(Exception):
 
 @dataclass(frozen=True)
 class AgentConfiguration:
+    """
+    A configuration for Infection Monkey agents
+
+    Attributes:
+        :param keep_tunnel_open_time: Maximum time in seconds to keep a tunnel open after
+                                      the last exploit
+        :param custom_pbas: Configuration for custom post-breach actions
+        :param post_breach_actions: Configuration for post-breach actions
+        :param credential_collectors: Configuration for credential collectors
+        :param payloads: Configuration for payloads
+        :param propagation: Configuration for propagation
+    """
+
     keep_tunnel_open_time: float
     custom_pbas: CustomPBAConfiguration
     post_breach_actions: Tuple[PluginConfiguration, ...]
@@ -104,7 +117,7 @@ class AgentConfiguration:
 
 
 class AgentConfigurationSchema(Schema):
-    keep_tunnel_open_time = fields.Float()
+    keep_tunnel_open_time = fields.Float(validate=validate.Range(min=0))
     custom_pbas = fields.Nested(CustomPBAConfigurationSchema)
     post_breach_actions = fields.List(fields.Nested(PluginConfigurationSchema))
     credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema))

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

@@ -251,6 +251,14 @@ def test_agent_configuration():
     assert json.loads(config_json) == AGENT_CONFIGURATION
 
 
+def test_agent_configuration__negative_keep_tunnel_open_time():
+    negative_keep_tunnel_open_time_configuration = AGENT_CONFIGURATION.copy()
+    negative_keep_tunnel_open_time_configuration["keep_tunnel_open_time"] = -1
+
+    with pytest.raises(InvalidConfigurationError):
+        AgentConfiguration.from_mapping(negative_keep_tunnel_open_time_configuration)
+
+
 def test_incorrect_type():
     valid_config = AgentConfiguration.from_mapping(AGENT_CONFIGURATION)
     with pytest.raises(InvalidConfigurationError):