diff --git a/monkey/monkey_island/cc/models/zero_trust/finding.py b/monkey/monkey_island/cc/models/zero_trust/finding.py index 2e6b9c172..5c8223e49 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/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 diff --git a/monkey/monkey_island/cc/models/zero_trust/monkey_finding.py b/monkey/monkey_island/cc/models/zero_trust/monkey_finding.py index 1de81303c..479b9b244 100644 --- a/monkey/monkey_island/cc/models/zero_trust/monkey_finding.py +++ b/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 diff --git a/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py b/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py index dd323a73c..9e36e46c5 100644 --- a/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py +++ b/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 diff --git a/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_finding_service.py b/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_finding_service.py index 72e9f7861..d8e439c71 100644 --- a/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_finding_service.py +++ b/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: