Common: Add CustomPBAConfiguration

monkey/common/configuration/__init__.py

@@ -1 +1,6 @@
-from .agent_configuration import PluginConfiguration, PluginConfigurationSchema
+from .agent_configuration import (
+    PluginConfiguration,
+    PluginConfigurationSchema,
+    CustomPBAConfiguration,
+    CustomPBAConfigurationSchema,
+)

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

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

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