make sure the cve enrichment skeleton works

Author: pyth0n1c

contentctl/enrichments/cve_enrichment.py

@@ -55,11 +55,11 @@ def getCveEnrichment(config:validate, timeout_seconds:int=10, force_disable_enri
         return CveEnrichment(use_enrichment=False, cve_api_obj=None)
 
 
-    def enrich_cve(self, cve_id:str, raise_exception_on_failure:bool=True)->Union[CveEnrichmentObj,None]:
+    def enrich_cve(self, cve_id:str, raise_exception_on_failure:bool=True)->CveEnrichmentObj:
 
         if not self.use_enrichment:
-            CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
+            return CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
         else:
             print("WARNING - Dynamic enrichment not supported at this time.")
-            CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
+            return CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
         # Depending on needs, we may add dynamic enrichment functionality back to the tool

contentctl/objects/abstract_security_content_objects/detection_abstract.py

@@ -146,7 +146,7 @@ def mappings(self)->dict[str, List[str]]:
     cve_enrichment: list[CveEnrichmentObj] = Field([], validate_default=True)
 
     @model_validator(mode="after")
-    def cve_enrichment_func(self, info:ValidationInfo)->Detection_Abstract:
+    def cve_enrichment_func(self, info:ValidationInfo):
         if len(self.cve_enrichment) > 0:
             raise ValueError(f"Error, field 'cve_enrichment' should be empty and "
                              f"dynamically populated at runtime. Instead, this field contained: {self.cve_enrichment}")
@@ -155,20 +155,15 @@ def cve_enrichment_func(self, info:ValidationInfo)->Detection_Abstract:
         if output_dto is None:
             raise ValueError("Context not provided to detection model post validator")
 
-        if output_dto.cve_enrichment.use_enrichment is False:    
-            return self
 
         enriched_cves:list[CveEnrichmentObj] = []
+
         for cve_id in self.tags.cve:
             try:
-                enrichment = output_dto.cve_enrichment.enrich_cve(cve_id, raise_exception_on_failure=False)
-                if enrichment is None:
-                    print(f"WARNING: Failed to find cve_id '{cve_id}'")
-                else:
-                    enriched_cves.append(enrichment)
+                enriched_cves.append(output_dto.cve_enrichment.enrich_cve(cve_id, raise_exception_on_failure=False))
             except Exception as e:
                 raise ValueError(f"{e}")
-
+        self.cve_enrichment = enriched_cves
         return self