From e0ae109368321fec5635cb1dee05d621676e3887 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Jun 2022 09:27:51 -0400 Subject: [PATCH] Common: Add AgentConfiguration --- monkey/common/configuration/__init__.py | 2 + .../configuration/agent_configuration.py | 23 ++++++ .../common/test_agent_configuration.py | 73 +++++++++++++------ vulture_allowlist.py | 1 + 4 files changed, 78 insertions(+), 21 deletions(-) diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index d49080618..973e6e28b 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -19,4 +19,6 @@ from .agent_configuration import ( ExploitationConfigurationSchema, PropagationConfiguration, PropagationConfigurationSchema, + AgentConfiguration, + AgentConfigurationSchema, ) diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py index 222ff707c..40e82d5c3 100644 --- a/monkey/common/configuration/agent_configuration.py +++ b/monkey/common/configuration/agent_configuration.py @@ -169,3 +169,26 @@ class PropagationConfigurationSchema(Schema): @post_load def _make_propagation_configuration(self, data, **kwargs): return PropagationConfiguration(**data) + + +@dataclass(frozen=True) +class AgentConfiguration: + keep_tunnel_open_time: float + custom_pbas: CustomPBAConfiguration + post_breach_actions: List[PluginConfiguration] + credential_collectors: List[PluginConfiguration] + payloads: List[PluginConfiguration] + propagation: PropagationConfiguration + + +class AgentConfigurationSchema(Schema): + keep_tunnel_open_time = fields.Float() + custom_pbas = fields.Nested(CustomPBAConfigurationSchema) + post_breach_actions = fields.List(fields.Nested(PluginConfigurationSchema)) + credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema)) + payloads = fields.List(fields.Nested(PluginConfigurationSchema)) + propagation = fields.Nested(PropagationConfigurationSchema) + + @post_load + def _make_agent_configuration(self, data, **kwargs): + return AgentConfiguration(**data) diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py index da52ab12a..6954407bd 100644 --- a/monkey/tests/unit_tests/common/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/test_agent_configuration.py @@ -1,5 +1,8 @@ from common import OperatingSystems from common.configuration import ( + AgentConfiguration, + AgentConfigurationSchema, + CustomPBAConfiguration, CustomPBAConfigurationSchema, ExploitationConfiguration, ExploitationConfigurationSchema, @@ -8,6 +11,7 @@ from common.configuration import ( ICMPScanConfigurationSchema, NetworkScanConfiguration, NetworkScanConfigurationSchema, + PluginConfiguration, PluginConfigurationSchema, PropagationConfiguration, PropagationConfigurationSchema, @@ -15,38 +19,41 @@ from common.configuration import ( TCPScanConfigurationSchema, ) +NAME = "bond" +OPTIONS = {"gun": "Walther PPK", "car": "Aston Martin DB5"} +PLUGIN_CONFIGURATION = {"name": NAME, "options": OPTIONS} + def test_build_plugin_configuration(): - name = "bond" - options = {"gun": "Walther PPK", "car": "Aston Martin DB5"} schema = PluginConfigurationSchema() - config = schema.load({"name": name, "options": options}) + config = schema.load(PLUGIN_CONFIGURATION) - assert config.name == name - assert config.options == options + assert config.name == NAME + assert config.options == OPTIONS + + +LINUX_COMMAND = "a" +LINUX_FILENAME = "b" +WINDOWS_COMMAND = "c" +WINDOWS_FILENAME = "d" +CUSTOM_PBA_CONFIGURATION = { + "linux_command": LINUX_COMMAND, + "linux_filename": LINUX_FILENAME, + "windows_command": WINDOWS_COMMAND, + "windows_filename": WINDOWS_FILENAME, +} 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, - } - ) + config = schema.load(CUSTOM_PBA_CONFIGURATION) - assert config.linux_command == linux_command - assert config.linux_filename == linux_filename - assert config.windows_command == windows_command - assert config.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 BLOCKED_IPS = ["10.0.0.1", "192.168.1.1"] @@ -197,3 +204,27 @@ def test_propagation_configuration(): assert isinstance(config.exploitation, ExploitationConfiguration) assert config.maximum_depth == 5 assert config_dict == PROPAGATION_CONFIGURATION + + +def test_agent_configuration(): + agent_configuration = { + "keep_tunnel_open_time": 30, + "custom_pbas": CUSTOM_PBA_CONFIGURATION, + "post_breach_actions": [PLUGIN_CONFIGURATION], + "credential_collectors": [PLUGIN_CONFIGURATION], + "payloads": [PLUGIN_CONFIGURATION], + "propagation": PROPAGATION_CONFIGURATION, + } + schema = AgentConfigurationSchema() + + config = schema.load(agent_configuration) + config_dict = schema.dump(config) + + assert isinstance(config, AgentConfiguration) + assert config.keep_tunnel_open_time == 30 + assert isinstance(config.custom_pbas, CustomPBAConfiguration) + assert isinstance(config.post_breach_actions[0], PluginConfiguration) + assert isinstance(config.credential_collectors[0], PluginConfiguration) + assert isinstance(config.payloads[0], PluginConfiguration) + assert isinstance(config.propagation, PropagationConfiguration) + assert config_dict == agent_configuration diff --git a/vulture_allowlist.py b/vulture_allowlist.py index a05814c54..3a8ad98c1 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -193,6 +193,7 @@ _make_icmp_scan_configuration # unused method (monkey/common/configuration/agen _make_tcp_scan_configuration # unused method (monkey/common/configuration/agent_configuration.py:122) _make_network_scan_configuration # unused method (monkey/common/configuration/agent_configuration.py:110) _make_propagation_configuration # unused method (monkey/common/configuration/agent_configuration.py:167) +_make_agent_configuration # unused method (monkey/common/configuration/agent_configuration.py:192) LINUX # unused variable (monkey/common/operating_systems.py:5) WINDOWS # unused variable (monkey/common/operating_systems.py:6)