diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index 9b5cdfb64..c36144499 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -9,6 +9,8 @@ from .agent_configuration import ( ExploiterConfigurationSchema, ExploitationConfiguration, ExploitationConfigurationSchema, + ICMPScanConfiguration, + ICMPScanConfigurationSchema, TCPScanConfiguration, TCPScanConfigurationSchema, ) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index 42f7dfcac..550d58d01 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -101,6 +101,14 @@ class ICMPScanConfiguration: timeout_ms: int +class ICMPScanConfigurationSchema(Schema): + timeout_ms = fields.Int() + + @post_load + def make_icmp_scan_configuration(self, data, **kwargs): + return ICMPScanConfiguration(**data) + + @dataclass(frozen=True) class TCPScanConfiguration: timeout_ms: int diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py index ad4e52790..d6c1a3a18 100644 --- a/monkey/tests/unit_tests/common/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/test_agent_configuration.py @@ -5,6 +5,7 @@ from common.configuration import ( ExploitationConfigurationSchema, ExploitationOptionsConfigurationSchema, ExploiterConfigurationSchema, + ICMPScanConfigurationSchema, PluginConfigurationSchema, TCPScanConfigurationSchema, ) @@ -98,6 +99,15 @@ def test_exploitation_configuration(): assert config_dict == exploitation_config +def test_icmp_scan_configuration_schema(): + timeout_ms = 2525 + schema = ICMPScanConfigurationSchema() + + config = schema.load({"timeout_ms": timeout_ms}) + + assert config.timeout_ms == timeout_ms + + def test_tcp_scan_configuration_schema(): timeout_ms = 2525 ports = [8080, 443]