Refactored details out of abstract finding document.

monkey/monkey_island/cc/models/zero_trust/finding.py

@@ -5,7 +5,7 @@ Define a Document Schema for Zero Trust findings.
 
 from __future__ import annotations
 
-from mongoengine import Document, GenericLazyReferenceField, StringField
+from mongoengine import Document, StringField
 
 import common.common_consts.zero_trust_consts as zero_trust_consts
 # Dummy import for mongoengine.
@@ -37,17 +37,3 @@ class Finding(Document):
     # SCHEMA
     test = StringField(required=True, choices=zero_trust_consts.TESTS)
     status = StringField(required=True, choices=zero_trust_consts.ORDERED_TEST_STATUSES)
-
-    # Details are in a separate document in order to discourage pulling them when not needed
-    # due to performance.
-    details = GenericLazyReferenceField(required=True)
-
-    # Creation methods
-    @classmethod
-    def save_finding(cls,
-                     test: str,
-                     status: str,
-                     detail_ref) -> Finding:
-        finding = cls(test=test, status=status, details=detail_ref)
-        finding.save()
-        return finding

monkey/monkey_island/cc/models/zero_trust/monkey_finding.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from mongoengine import LazyReferenceField
 
 from monkey_island.cc.models.zero_trust.finding import Finding
@@ -5,4 +7,14 @@ from monkey_island.cc.models.zero_trust.monkey_finding_details import MonkeyFind
 
 
 class MonkeyFinding(Finding):
+    # We put additional info into a lazy reference field, because this info should be only
+    # pulled when explicitly needed due to performance
     details = LazyReferenceField(MonkeyFindingDetails, required=True)
+
+    @staticmethod
+    def save_finding(test: str,
+                     status: str,
+                     detail_ref: MonkeyFindingDetails) -> MonkeyFinding:
+        finding = MonkeyFinding(test=test, status=status, details=detail_ref)
+        finding.save()
+        return finding

monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from mongoengine import LazyReferenceField
 
 from monkey_island.cc.models.zero_trust.finding import Finding
@@ -5,4 +7,14 @@ from monkey_island.cc.models.zero_trust.scoutsuite_finding_details import ScoutS
 
 
 class ScoutSuiteFinding(Finding):
+    # We put additional info into a lazy reference field, because this info should be only
+    # pulled when explicitly needed due to performance
     details = LazyReferenceField(ScoutSuiteFindingDetails, required=True)
+
+    @staticmethod
+    def save_finding(test: str,
+                     status: str,
+                     detail_ref: ScoutSuiteFindingDetails) -> ScoutSuiteFinding:
+        finding = ScoutSuiteFinding(test=test, status=status, details=detail_ref)
+        finding.save()
+        return finding

monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_finding_service.py

@@ -19,7 +19,7 @@ class MonkeyZTFindingService:
         :raises: Assertion error if this is used when there's more then one finding which fits the query - this is not
         when this function should be used.
         """
-        existing_findings = MonkeyFinding.objects(test=test, status=status)
+        existing_findings = list(MonkeyFinding.objects(test=test, status=status))
         assert (len(existing_findings) < 2), "More than one finding exists for {}:{}".format(test, status)
 
         if len(existing_findings) == 0: