diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index a40b7a97e..8d71636a4 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -1 +1,6 @@ -from .agent_configuration import PluginConfiguration, PluginConfigurationSchema +from .agent_configuration import ( + PluginConfiguration, + PluginConfigurationSchema, + CustomPBAConfiguration, + CustomPBAConfigurationSchema, +) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index fbcc4872d..f89e27e69 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -3,6 +3,28 @@ from dataclasses import dataclass from marshmallow import RAISE, Schema, fields, post_load +@dataclass(frozen=True) +class CustomPBAConfiguration: + linux_command: str + linux_filename: str + windows_command: str + windows_filename: str + + +class CustomPBAConfigurationSchema(Schema): + class Meta: + unknown = RAISE + + linux_command = fields.Str() + linux_filename = fields.Str() + windows_command = fields.Str() + windows_filename = fields.Str() + + @post_load + def make_custom_pba_configuration(self, data, **kwargs): + return CustomPBAConfiguration(**data) + + @dataclass(frozen=True) class PluginConfiguration: name: str diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py index 53555680d..38c744b17 100644 --- a/monkey/tests/unit_tests/common/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/test_agent_configuration.py @@ -1,12 +1,34 @@ -from common.configuration import PluginConfigurationSchema +from common.configuration import CustomPBAConfigurationSchema, PluginConfigurationSchema def test_build_plugin_configuration(): name = "bond" options = {"gun": "Walther PPK", "car": "Aston Martin DB5"} - pcs = PluginConfigurationSchema() + schema = PluginConfigurationSchema() - pc = pcs.load({"name": name, "options": options}) + config = schema.load({"name": name, "options": options}) - assert pc.name == name - assert pc.options == options + assert config.name == name + assert config.options == options + + +def test_custom_pba_configuration_schema(): + linux_command = "a" + linux_filename = "b" + windows_command = "c" + windows_filename = "d" + schema = CustomPBAConfigurationSchema() + + config = schema.load( + { + "linux_command": linux_command, + "linux_filename": linux_filename, + "windows_command": windows_command, + "windows_filename": windows_filename, + } + ) + + assert config.linux_command == linux_command + assert config.linux_filename == linux_filename + assert config.windows_command == windows_command + assert config.windows_filename == windows_filename diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 73961d4d7..9e3522f86 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -187,6 +187,7 @@ response_code # unused variable (monkey/monkey_island/cc/services/aws/aws_comma Meta # unused class(monkey/common/configuration/agent_configuration.py:13) unknown # unused variable(monkey/common/configuration/agent_configuration.py:14) make_plugin_configuration # unused method(monkey/common/configuration/agent_configuration.py:19) +make_custom_pba_configuration # unused method(monkey/common/configuration/agent_configuration.py:34) # TODO DELETE AFTER RESOURCE REFACTORING