From 4065bc23fb61ad5d90fb3cb7c90f35242d2ee2a8 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 16 Jun 2022 20:17:22 -0400
Subject: [PATCH] Common: Add ExploitationConfiguration

---
 monkey/common/configuration/__init__.py       |  2 ++
 .../configuration/agent_configuration.py      | 17 ++++++++++
 .../common/test_agent_configuration.py        | 33 +++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py
index f9b876cdf..c22370603 100644
--- a/monkey/common/configuration/__init__.py
+++ b/monkey/common/configuration/__init__.py
@@ -7,4 +7,6 @@ from .agent_configuration import (
     ExploitationOptionsConfigurationSchema,
     ExploiterConfiguration,
     ExploiterConfigurationSchema,
+    ExploitationConfiguration,
+    ExploitationConfigurationSchema,
 )
diff --git a/monkey/common/configuration/agent_configuration.py b/monkey/common/configuration/agent_configuration.py
index bb7e08bbc..e961a82f3 100644
--- a/monkey/common/configuration/agent_configuration.py
+++ b/monkey/common/configuration/agent_configuration.py
@@ -69,3 +69,20 @@ class ExploiterConfigurationSchema(Schema):
     @post_load
     def make_exploiter_configuration(self, data, **kwargs):
         return ExploiterConfiguration(**data)
+
+
+@dataclass(frozen=True)
+class ExploitationConfiguration:
+    options: ExploitationOptionsConfiguration
+    brute_force: List[ExploiterConfiguration]
+    vulnerability: List[ExploiterConfiguration]
+
+
+class ExploitationConfigurationSchema(Schema):
+    options = fields.Nested(ExploitationOptionsConfigurationSchema)
+    brute_force = fields.List(fields.Nested(ExploiterConfigurationSchema))
+    vulnerability = fields.List(fields.Nested(ExploiterConfigurationSchema))
+
+    @post_load
+    def make_exploitation_options_configuration(self, data, **kwargs):
+        return ExploitationConfiguration(**data)
diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py
index 21c0c8575..cee26fe5b 100644
--- a/monkey/tests/unit_tests/common/test_agent_configuration.py
+++ b/monkey/tests/unit_tests/common/test_agent_configuration.py
@@ -1,6 +1,8 @@
 from common import OperatingSystems
 from common.configuration import (
     CustomPBAConfigurationSchema,
+    ExploitationConfiguration,
+    ExploitationConfigurationSchema,
     ExploitationOptionsConfigurationSchema,
     ExploiterConfigurationSchema,
     PluginConfigurationSchema,
@@ -62,3 +64,34 @@ def test_exploiter_configuration_schema():
     assert config.name == name
     assert config.options == options
     assert config.supported_os == supported_os
+
+
+def test_exploitation_configuration():
+    ports = [1, 2, 3]
+    brute_force = [
+        {"name": "ex1", "options": {}, "supported_os": ["LINUX"]},
+        {
+            "name": "ex2",
+            "options": {"smb_download_timeout": 10},
+            "supported_os": ["LINUX", "WINDOWS"],
+        },
+    ]
+    vulnerability = [
+        {
+            "name": "ex3",
+            "options": {"smb_download_timeout": 10},
+            "supported_os": ["WINDOWS"],
+        },
+    ]
+    exploitation_config = {
+        "options": {"http_ports": ports},
+        "brute_force": brute_force,
+        "vulnerability": vulnerability,
+    }
+    schema = ExploitationConfigurationSchema()
+
+    config = schema.load(exploitation_config)
+    config_dict = schema.dump(config)
+
+    assert isinstance(config, ExploitationConfiguration)
+    assert config_dict == exploitation_config