From 85a11b62bde23c9848757966ab622bf6bcb0b490 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 27 Jul 2022 18:26:32 +0530 Subject: [PATCH 1/3] Common: Add docstring to AgentConfiguration dataclass --- .../agent_configuration/agent_configuration.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/monkey/common/agent_configuration/agent_configuration.py b/monkey/common/agent_configuration/agent_configuration.py index 00b39dce4..0b45c5c5f 100644 --- a/monkey/common/agent_configuration/agent_configuration.py +++ b/monkey/common/agent_configuration/agent_configuration.py @@ -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, ...] From a603b16573c2fbb9f0e4d236b47213c281f0580d Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 27 Jul 2022 18:28:52 +0530 Subject: [PATCH 2/3] Common: Add maximum tunnel open time validation to AgentConfigurationSchema --- monkey/common/agent_configuration/agent_configuration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/common/agent_configuration/agent_configuration.py b/monkey/common/agent_configuration/agent_configuration.py index 0b45c5c5f..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 @@ -117,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)) From e6108d0fc77da9daef2432f7e213f7ebf65f11d2 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 27 Jul 2022 18:32:02 +0530 Subject: [PATCH 3/3] UT: Add test for AgentConfigurationSchema tunnel open time validation --- .../common/configuration/test_agent_configuration.py | 8 ++++++++ 1 file changed, 8 insertions(+) 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 95539b9c7..0d6ffa3c6 100644 --- a/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py @@ -239,6 +239,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):