diff --git a/monkey/monkey_island/cc/models/base_models.py b/monkey/monkey_island/cc/models/base_models.py new file mode 100644 index 000000000..5d82be324 --- /dev/null +++ b/monkey/monkey_island/cc/models/base_models.py @@ -0,0 +1,51 @@ +import json +from typing import Sequence + +from pydantic import BaseModel, Extra, ValidationError + + +class InfectionMonkeyModelConfig: + underscore_attrs_are_private = True + extra = Extra.forbid + + +class InfectionMonkeyBaseModel(BaseModel): + class Config(InfectionMonkeyModelConfig): + pass + + def __init__(self, **kwargs): + try: + super().__init__(**kwargs) + except ValidationError as err: + # TLDR: This exception handler allows users of this class to be decoupled from pydantic. + # + # When validation of a pydantic object fails, pydantic raises a `ValidationError`, which + # is a `ValueError`, even if the real cause was a `TypeError`. Furthermore, allowing + # `pydantic.ValueError` to be raised would couple other modules to pydantic, which is + # undesirable. This exception handler re-raises the first validation error that pydantic + # encountered. This allows users of these models to `except` `TypeError` or `ValueError` + # and handle them. Pydantic-specific errors are still raised, but they inherit from + # `TypeError` or `ValueError`. + e = err.raw_errors[0] + while isinstance(e, Sequence): + e = e[0] + + raise e.exc + + # We need to be able to convert our models to fully simplified dictionaries. The + # `BaseModel.dict()` does not support this. There is a proposal to add a `simplify` keyword + # argument to `dict()` to support this. See + # https://github.com/pydantic/pydantic/issues/951#issuecomment-552463606. The hope is that we + # can override `dict()` with an implementation of `simplify` and remove it when the feature gets + # merged. If the feature doesn't get merged, or the interface is changed, this function can + # continue to serve as a wrapper until we can update all references to it. + def dict(self, simplify=False, **kwargs): + if simplify: + return json.loads(self.json()) + return BaseModel.dict(self, **kwargs) + + +class MutableBaseModel(InfectionMonkeyBaseModel): + class Config(InfectionMonkeyModelConfig): + allow_mutation = True + validate_assignment = True diff --git a/vulture_allowlist.py b/vulture_allowlist.py index b340dc97c..70f620492 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -274,3 +274,9 @@ serialize event deserialize serialized_event + +# pydantic base models +underscore_attrs_are_private +extra +allow_mutation +validate_assignment