From 09474ac1fe0de79ccaee7abbab545f4e541784a5 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 18 Aug 2022 08:44:46 -0400
Subject: [PATCH] Island: Add base models for pydantic classes

---
 monkey/monkey_island/cc/models/base_models.py | 51 +++++++++++++++++++
 vulture_allowlist.py                          |  6 +++
 2 files changed, 57 insertions(+)
 create mode 100644 monkey/monkey_island/cc/models/base_models.py

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