Common: Remove old AgentConfiguration using marshmallow

This commit is contained in:
Shreya Malviya 2022-08-29 17:20:06 +05:30
parent 6d29829808
commit 1311fd5d27
1 changed files with 1 additions and 108 deletions

View File

@ -1,24 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Mapping, Tuple
from typing import Tuple
from marshmallow import Schema, fields, validate
from marshmallow.exceptions import MarshmallowError
from pydantic import PositiveFloat
from common.base_models import MutableInfectionMonkeyBaseModel
from ..utils.code_utils import freeze_lists_in_mapping
from .agent_sub_configuration_schemas import (
CustomPBAConfigurationSchema,
PluginConfigurationSchema,
PropagationConfigurationSchema,
)
from .agent_sub_configurations import (
CustomPBAConfiguration,
PluginConfiguration,
PropagationConfiguration,
Pydantic___CustomPBAConfiguration,
Pydantic___PluginConfiguration,
Pydantic___PropagationConfiguration,
@ -36,101 +24,6 @@ class InvalidConfigurationError(Exception):
)
@dataclass(frozen=True)
class AgentConfiguration:
"""
A configuration for Infection Monkey agents
Attributes:
:param keep_tunnel_open_time: Maximum time in seconds to keep a tunnel open after
the last exploit
:param custom_pbas: Configuration for custom post-breach actions
:param post_breach_actions: Configuration for post-breach actions
:param credential_collectors: Configuration for credential collectors
:param payloads: Configuration for payloads
:param propagation: Configuration for propagation
"""
keep_tunnel_open_time: float
custom_pbas: CustomPBAConfiguration
post_breach_actions: Tuple[PluginConfiguration, ...]
credential_collectors: Tuple[PluginConfiguration, ...]
payloads: Tuple[PluginConfiguration, ...]
propagation: PropagationConfiguration
def __post_init__(self):
# This will raise an exception if the object is invalid. Calling this in __post__init()
# makes it impossible to construct an invalid object
try:
AgentConfigurationSchema().dump(self)
except Exception as err:
raise InvalidConfigurationError(str(err))
@staticmethod
def from_mapping(config_mapping: Mapping[str, Any]) -> AgentConfiguration:
"""
Construct an AgentConfiguration from a Mapping
:param config_mapping: A Mapping that represents an AgentConfiguration
:return: An AgentConfiguration
:raises: InvalidConfigurationError if the provided Mapping does not represent a valid
AgentConfiguration
"""
try:
config_dict = AgentConfigurationSchema().load(config_mapping)
config_dict = freeze_lists_in_mapping(config_dict)
return AgentConfiguration(**config_dict)
except MarshmallowError as err:
raise InvalidConfigurationError(str(err))
@staticmethod
def from_json(config_json: str) -> AgentConfiguration:
"""
Construct an AgentConfiguration from a JSON string
:param config_json: A JSON string that represents an AgentConfiguration
:return: An AgentConfiguration
:raises: InvalidConfigurationError if the provided JSON does not represent a valid
AgentConfiguration
"""
try:
config_dict = AgentConfigurationSchema().loads(config_json)
config_dict = freeze_lists_in_mapping(config_dict)
return AgentConfiguration(**config_dict)
except MarshmallowError as err:
raise InvalidConfigurationError(str(err))
@staticmethod
def to_mapping(config: AgentConfiguration) -> Mapping[str, Any]:
"""
Serialize an AgentConfiguration to a Mapping
:param config: An AgentConfiguration
:return: A Mapping that represents the AgentConfiguration
"""
return AgentConfigurationSchema().dump(config)
@staticmethod
def to_json(config: AgentConfiguration) -> str:
"""
Serialize an AgentConfiguration to JSON
:param config: An AgentConfiguration
:return: A JSON string that represents the AgentConfiguration
"""
return AgentConfigurationSchema().dumps(config)
class AgentConfigurationSchema(Schema):
keep_tunnel_open_time = fields.Float(validate=validate.Range(min=0))
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)
class Pydantic___AgentConfiguration(MutableInfectionMonkeyBaseModel):
keep_tunnel_open_time: PositiveFloat
custom_pbas: Pydantic___CustomPBAConfiguration