Common: Use OperatingSystems enum in ExploiterConfigurationSchema

This commit is contained in:
Mike Salvatore 2022-06-16 19:52:12 -04:00
parent 737aacf0bd
commit 70e8bca1ea
3 changed files with 17 additions and 9 deletions

View File

@ -2,6 +2,9 @@ from dataclasses import dataclass
from typing import Dict, List from typing import Dict, List
from marshmallow import RAISE, Schema, fields, post_load from marshmallow import RAISE, Schema, fields, post_load
from marshmallow_enum import EnumField
from common import OperatingSystems
@dataclass(frozen=True) @dataclass(frozen=True)
@ -48,7 +51,7 @@ class PluginConfigurationSchema(Schema):
class ExploiterConfiguration: class ExploiterConfiguration:
name: str name: str
options: Dict options: Dict
supported_os: List[str] supported_os: List[OperatingSystems]
class ExploiterConfigurationSchema(Schema): class ExploiterConfigurationSchema(Schema):
@ -57,7 +60,7 @@ class ExploiterConfigurationSchema(Schema):
name = fields.Str() name = fields.Str()
options = fields.Mapping() options = fields.Mapping()
supported_os = fields.List(fields.Str()) supported_os = fields.List(EnumField(OperatingSystems))
@post_load @post_load
def make_exploiter_configuration(self, data, **kwargs): def make_exploiter_configuration(self, data, **kwargs):

View File

@ -1,3 +1,4 @@
from common import OperatingSystems
from common.configuration import ( from common.configuration import (
CustomPBAConfigurationSchema, CustomPBAConfigurationSchema,
ExploiterConfigurationSchema, ExploiterConfigurationSchema,
@ -41,10 +42,12 @@ def test_custom_pba_configuration_schema():
def test_exploiter_configuration_schema(): def test_exploiter_configuration_schema():
name = "bond" name = "bond"
options = {"gun": "Walther PPK", "car": "Aston Martin DB5"} options = {"gun": "Walther PPK", "car": "Aston Martin DB5"}
supported_os = ["linux", "windows"] supported_os = [OperatingSystems.LINUX, OperatingSystems.WINDOWS]
schema = ExploiterConfigurationSchema() schema = ExploiterConfigurationSchema()
config = schema.load({"name": name, "options": options, "supported_os": supported_os}) config = schema.load(
{"name": name, "options": options, "supported_os": [os_.name for os_ in supported_os]}
)
assert config.name == name assert config.name == name
assert config.options == options assert config.options == options

View File

@ -184,11 +184,13 @@ architecture # unused variable (monkey/infection_monkey/exploit/caching_agent_r
response_code # unused variable (monkey/monkey_island/cc/services/aws/aws_command_runner.py:26) response_code # unused variable (monkey/monkey_island/cc/services/aws/aws_command_runner.py:26)
# Agent Configuration # Agent Configuration
Meta # unused class(monkey/common/configuration/agent_configuration.py:13) Meta # unused class (monkey/common/configuration/agent_configuration.py:13)
unknown # unused variable(monkey/common/configuration/agent_configuration.py:14) unknown # unused variable (monkey/common/configuration/agent_configuration.py:14)
make_plugin_configuration # unused method(monkey/common/configuration/agent_configuration.py:19) 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_custom_pba_configuration # unused method (monkey/common/configuration/agent_configuration.py:34)
make_exploiter_configuration # unused method(monkey/common/configuration/agent_configuration.py:62) make_exploiter_configuration # unused method (monkey/common/configuration/agent_configuration.py:62)
LINUX # unused variable (monkey/common/operating_systems.py:5)
WINDOWS # unused variable (monkey/common/operating_systems.py:6)
# TODO DELETE AFTER RESOURCE REFACTORING # TODO DELETE AFTER RESOURCE REFACTORING