From 1311fd5d27e16acf8e2a1d46f9d7ba152778e742 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 29 Aug 2022 17:20:06 +0530 Subject: [PATCH] Common: Remove old AgentConfiguration using marshmallow --- .../agent_configuration.py | 109 +----------------- 1 file changed, 1 insertion(+), 108 deletions(-) diff --git a/monkey/common/agent_configuration/agent_configuration.py b/monkey/common/agent_configuration/agent_configuration.py index 84782bef2..34635577f 100644 --- a/monkey/common/agent_configuration/agent_configuration.py +++ b/monkey/common/agent_configuration/agent_configuration.py @@ -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