diff --git a/monkey/common/agent_configuration/agent_configuration.py b/monkey/common/agent_configuration/agent_configuration.py index 00b39dce4..cfd685d0f 100644 --- a/monkey/common/agent_configuration/agent_configuration.py +++ b/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)) diff --git a/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py b/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py index ec1f884e1..6110dbc7b 100644 --- a/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py +++ b/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):