Common: Create AgentConfiguration using pydantic

This commit is contained in:
Shreya Malviya 2022-08-29 13:58:44 +05:30
parent b8914101a6
commit 2f05d22780
1 changed files with 15 additions and 0 deletions

View File

@ -5,6 +5,9 @@ from typing import Any, Mapping, Tuple
from marshmallow import Schema, fields, validate from marshmallow import Schema, fields, validate
from marshmallow.exceptions import MarshmallowError 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 ..utils.code_utils import freeze_lists_in_mapping
from .agent_sub_configuration_schemas import ( from .agent_sub_configuration_schemas import (
@ -16,6 +19,9 @@ from .agent_sub_configurations import (
CustomPBAConfiguration, CustomPBAConfiguration,
PluginConfiguration, PluginConfiguration,
PropagationConfiguration, PropagationConfiguration,
Pydantic___CustomPBAConfiguration,
Pydantic___PluginConfiguration,
Pydantic___PropagationConfiguration,
) )
@ -123,3 +129,12 @@ class AgentConfigurationSchema(Schema):
credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema)) credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema))
payloads = fields.List(fields.Nested(PluginConfigurationSchema)) payloads = fields.List(fields.Nested(PluginConfigurationSchema))
propagation = fields.Nested(PropagationConfigurationSchema) propagation = fields.Nested(PropagationConfigurationSchema)
class Pydantic___AgentConfiguration(MutableInfectionMonkeyBaseModel):
keep_tunnel_open_time: PositiveFloat
custom_pbas: Pydantic___CustomPBAConfiguration
post_breach_actions: Tuple[Pydantic___PluginConfiguration, ...]
credential_collectors: Tuple[Pydantic___PluginConfiguration, ...]
payloads: Tuple[Pydantic___PluginConfiguration, ...]
propagation: Pydantic___PropagationConfiguration