forked from p15670423/monkey
Common: Remove old AgentConfiguration using marshmallow
This commit is contained in:
parent
6d29829808
commit
1311fd5d27
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue