From 9d73252ff5bb947b651e09c65090d297090b94c7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 16 Jun 2022 15:53:34 -0400 Subject: [PATCH] Common: Add ExploiterConfiguration --- monkey/common/configuration/__init__.py | 2 ++ .../configuration/agent_configuration.py | 22 ++++++++++++++++++- .../common/test_agent_configuration.py | 19 +++++++++++++++- vulture_allowlist.py | 1 + 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index 8d71636a4..ed2b3cbc9 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -3,4 +3,6 @@ from .agent_configuration import ( PluginConfigurationSchema, CustomPBAConfiguration, CustomPBAConfigurationSchema, + ExploiterConfiguration, + ExploiterConfigurationSchema, ) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index b52b7cbc0..95592ecb7 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Dict +from typing import Dict, List from marshmallow import RAISE, Schema, fields, post_load @@ -42,3 +42,23 @@ class PluginConfigurationSchema(Schema): @post_load def make_plugin_configuration(self, data, **kwargs): return PluginConfiguration(**data) + + +@dataclass(frozen=True) +class ExploiterConfiguration: + name: str + options: Dict + supported_os: List[str] + + +class ExploiterConfigurationSchema(Schema): + class Meta: + unknown = RAISE + + name = fields.Str() + options = fields.Mapping() + supported_os = fields.List(fields.Str()) + + @post_load + def make_exploiter_configuration(self, data, **kwargs): + return ExploiterConfiguration(**data) diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py index 38c744b17..7854a48ea 100644 --- a/monkey/tests/unit_tests/common/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/test_agent_configuration.py @@ -1,4 +1,8 @@ -from common.configuration import CustomPBAConfigurationSchema, PluginConfigurationSchema +from common.configuration import ( + CustomPBAConfigurationSchema, + ExploiterConfigurationSchema, + PluginConfigurationSchema, +) def test_build_plugin_configuration(): @@ -32,3 +36,16 @@ def test_custom_pba_configuration_schema(): assert config.linux_filename == linux_filename assert config.windows_command == windows_command assert config.windows_filename == windows_filename + + +def test_exploiter_configuration_schema(): + name = "bond" + options = {"gun": "Walther PPK", "car": "Aston Martin DB5"} + supported_os = ["linux", "windows"] + schema = ExploiterConfigurationSchema() + + config = schema.load({"name": name, "options": options, "supported_os": supported_os}) + + assert config.name == name + assert config.options == options + assert config.supported_os == supported_os diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 9e3522f86..dec1aa7dd 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -188,6 +188,7 @@ 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) +make_exploiter_configuration # unused method(monkey/common/configuration/agent_configuration.py:62) # TODO DELETE AFTER RESOURCE REFACTORING