Enable CVE Enrichment.

However, there are issues with the API
availability/stability that we must overcome.

Author: pyth0n1c

contentctl/actions/validate.py

@@ -22,7 +22,8 @@ def execute(self, input_dto: validate) -> DirectorOutputDto:
 
         director_output_dto = DirectorOutputDto(AtomicTest.getAtomicTestsFromArtRepo(repo_path=input_dto.getAtomicRedTeamRepoPath(), 
                                                                                      enabled=input_dto.enrichments),
-                                                AttackEnrichment.getAttackEnrichment(input_dto),
+                                                AttackEnrichment.getAttackEnrichment(input_dto,
+                                                CveEnrichment.getCveEnrichment(input_dto),
                                                 [],[],[],[],[],[],[],[],[])
 
 

contentctl/enrichments/cve_enrichment.py

@@ -4,67 +4,17 @@
 import os
 import shelve
 import time
-from typing import Annotated
-from pydantic import BaseModel,Field,ConfigDict
-
+from typing import Annotated, Any, Union, TYPE_CHECKING
+from pydantic import BaseModel,Field
 from decimal import Decimal
-CVESSEARCH_API_URL = 'https://cve.circl.lu'
-
-CVE_CACHE_FILENAME = "lookups/CVE_CACHE.db"
-
-NON_PERSISTENT_CACHE = {}
-
-
-''''''
-@functools.cache
-def cvesearch_helper(url:str, cve_id:str, force_cached_or_offline:bool=False, max_api_attempts:int=3, retry_sleep_seconds:int=5):
-    if max_api_attempts < 1:
-            raise(Exception(f"The minimum number of CVESearch API attempts is 1.  You have passed {max_api_attempts}"))
-
-    if force_cached_or_offline:
-        if not os.path.exists(CVE_CACHE_FILENAME):
-            print(f"Cache at {CVE_CACHE_FILENAME} not found - Creating it.")
-        cache = shelve.open(CVE_CACHE_FILENAME, flag='c', writeback=True)
-    else:
-        cache = NON_PERSISTENT_CACHE
-    if cve_id in cache:
-        result = cache[cve_id]
-        #print(f"hit cve_enrichment:  {time.time() - start:.2f}")
-    else:
-        api_attempts_remaining = max_api_attempts
-        result = None
-        while api_attempts_remaining > 0:
-            api_attempts_remaining -= 1
-            try:
-                cve = cvesearch_id_helper(url)
-                result = cve.id(cve_id)
-                break
-            except Exception as e:
-                if api_attempts_remaining > 0:
-                    print(f"The option 'force_cached_or_offline' was used, but {cve_id} not found in {CVE_CACHE_FILENAME} and unable to connect to {CVESSEARCH_API_URL}: {str(e)}")
-                    print(f"Retrying the CVESearch API up to {api_attempts_remaining} more times after a sleep of {retry_sleep_seconds} seconds...")
-                    time.sleep(retry_sleep_seconds)
-                else:
-                    raise(Exception(f"The option 'force_cached_or_offline' was used, but {cve_id} not found in {CVE_CACHE_FILENAME} and unable to connect to {CVESSEARCH_API_URL} after {max_api_attempts} attempts: {str(e)}"))
-            
-        if result is None:
-            raise(Exception(f'CveEnrichment for [ {cve_id} ] failed - CVE does not exist'))
-        cache[cve_id] = result
-        
-    if isinstance(cache, shelve.Shelf):
-        #close the cache if it was a shelf
-        cache.close()
+from requests.exceptions import ReadTimeout
 
-    return result
+if TYPE_CHECKING:
+    from contentctl.objects.config import validate
 
-@functools.cache
-def cvesearch_id_helper(url:str):
-    #The initial CVESearch call takes some time.
-    #We cache it to avoid making this call each time we need to do a lookup
-    cve = CVESearch(CVESSEARCH_API_URL)
-    return cve
 
 
+CVESSEARCH_API_URL = 'https://cve.circl.lu'
 
 
 class CveEnrichmentObj(BaseModel):
@@ -74,27 +24,74 @@ class CveEnrichmentObj(BaseModel):
 
 
     @staticmethod
-    def buildEnrichmentOnFailure(id:Annotated[str, "^CVE-[1|2][0-9]{3}-[0-9]+$"], errorMessage:str)->CveEnrichmentObj:
+    def buildEnrichmentOnFailure(id:Annotated[str, "^CVE-[1|2][0-9]{3}-[0-9]+$"], errorMessage:str, 
+                                 raise_exception_on_failure:bool=True)->CveEnrichmentObj:
+        if raise_exception_on_failure:
+            raise Exception(errorMessage)
         message = f"{errorMessage}. Default CVSS of 5.0 used"
         print(message)
         return CveEnrichmentObj(id=id, cvss=Decimal(5.0), summary=message)
 
-class CveEnrichment():
-    @classmethod
-    def enrich_cve(cls, cve_id: str, force_cached_or_offline: bool = False, treat_failures_as_warnings:bool=True) -> CveEnrichmentObj:
-        cve_enriched = dict()
+
+# We need a MUCH better way to handle issues with the cve.circl.lu API.
+# It is often extremely slow or down, which means that we cannot enrich CVEs.
+# Downloading the entire database is VERY large, but I don't know that there
+# is an alternative.
+# Being able to include CVEs that have not made it into this database, or additonal
+# enriching comments on pre-existing CVEs, would also be extremely useful.
+timeout_error = False
+class CveEnrichment(BaseModel):
+    use_enrichment: bool = True
+    cve_api_obj: Union[CVESearch,None] = None
+    
+
+    class Config:
+        # Arbitrary_types are allowed to let us use the CVESearch Object
+        arbitrary_types_allowed = True
+        frozen = True
+        
+
+    @staticmethod
+    def getCveEnrichment(config:validate, timeout_seconds:int=10)->CveEnrichment:
+
+        if config.enrichments:
+            try:
+                cve_api_obj = CVESearch(CVESSEARCH_API_URL, timeout=timeout_seconds)
+                return CveEnrichment(use_enrichment=True, cve_api_obj=cve_api_obj)
+            except Exception as e:
+                raise Exception(f"Error setting CVE_SEARCH API to: {CVESSEARCH_API_URL}: {str(e)}")
+        
+        return CveEnrichment(use_enrichment=False, cve_api_obj=None)
+
+
+    @functools.cache
+    def enrich_cve(self, cve_id:str, raise_exception_on_failure:bool=True)->Union[CveEnrichmentObj,None]:
+        global timeout_error
+
+        if not self.use_enrichment:
+            return None
+        
+        if timeout_error:
+            message = f"Previous timeout during enrichment - CVE {cve_id} enrichment skipped."
+            return CveEnrichmentObj.buildEnrichmentOnFailure(id = cve_id, errorMessage=f"WARNING, {message}", 
+                                                             raise_exception_on_failure=raise_exception_on_failure) 
+ 
+        cve_enriched:dict[str,Any] = dict()
+
         try:
-            
-            result = cvesearch_helper(CVESSEARCH_API_URL, cve_id, force_cached_or_offline)
+            result = self.cve_api_obj.id(cve_id)
             cve_enriched['id'] = cve_id
             cve_enriched['cvss'] = result['cvss']
             cve_enriched['summary'] = result['summary']
+            return CveEnrichmentObj.model_validate(cve_enriched)
+        except ReadTimeout as e:
+            message = f"Timeout enriching CVE {cve_id}: {str(e)} after {self.cve_api_obj.timeout} seconds."\
+                      f" All other CVE Enrichment has been disabled"
+            #Set a global value to true so future runs don't waste time on this
+            timeout_error = True
+            return CveEnrichmentObj.buildEnrichmentOnFailure(id = cve_id, errorMessage=f"ERROR, {message}", 
+                                                             raise_exception_on_failure=raise_exception_on_failure)
         except Exception as e:
-            message = f"issue enriching {cve_id}, with error: {str(e)}"
-            if treat_failures_as_warnings: 
-                return CveEnrichmentObj.buildEnrichmentOnFailure(id = cve_id, errorMessage=f"WARNING, {message}")
-            else:
-                raise ValueError(f"ERROR, {message}")
-        
-        return CveEnrichmentObj.model_validate(cve_enriched)
-        
+            message = f"Error enriching CVE {cve_id}. Are you positive this CVE exists: {str(e)}"
+            return CveEnrichmentObj.buildEnrichmentOnFailure(id = cve_id, errorMessage=f"WARNING, {message}", 
+                                                             raise_exception_on_failure=raise_exception_on_failure)

contentctl/input/director.py

@@ -35,6 +35,7 @@ class DirectorOutputDto:
      # is far quicker than attack_enrichment
      atomic_tests: Union[list[AtomicTest],None]
      attack_enrichment: AttackEnrichment
+     cve_enrichment: CveEnrichment
      detections: list[Detection]
      stories: list[Story]
      baselines: list[Baseline]
@@ -44,7 +45,7 @@ class DirectorOutputDto:
      lookups: list[Lookup]
      deployments: list[Deployment]
      ssa_detections: list[SSADetection]
-     #cve_enrichment: CveEnrichment
+     
 
      name_to_content_map: dict[str, SecurityContentObject] = field(default_factory=dict)
      uuid_to_content_map: dict[UUID, SecurityContentObject] = field(default_factory=dict)

contentctl/objects/abstract_security_content_objects/detection_abstract.py

@@ -26,7 +26,7 @@
 
 #from contentctl.objects.playbook import Playbook
 from contentctl.objects.enums import DataSource,ProvidingTechnology
-from contentctl.enrichments.cve_enrichment import CveEnrichment, CveEnrichmentObj
+from contentctl.enrichments.cve_enrichment import CveEnrichmentObj
 
 
 class Detection_Abstract(SecurityContentObject):
@@ -143,18 +143,35 @@ def mappings(self)->dict[str, List[str]]:
     macros: list[Macro] = Field([],validate_default=True)
     lookups: list[Lookup] = Field([],validate_default=True)
 
-    @computed_field
-    @property
-    def cve_enrichment(self)->List[CveEnrichmentObj]:
+    cve_enrichment: list[CveEnrichmentObj] = Field([], validate_default=True)
+    
+    @model_validator(mode="after")
+    def cve_enrichment_func(self, info:ValidationInfo)->Detection_Abstract:
+        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}")
+
+        output_dto:Union[DirectorOutputDto,None]= info.context.get("output_dto",None)
+        if output_dto is None:
+            raise ValueError("Context not provided to detection model post validator")
 
-        raise Exception("CVE Enrichment Functionality not currently supported.  It will be re-added at a later time.")
-        enriched_cves = []
+        if output_dto.cve_enrichment.use_enrichment is False:    
+            return self
+        
+        enriched_cves:list[CveEnrichmentObj] = []
         for cve_id in self.tags.cve:
-            print(f"\nEnriching {cve_id}\n")
-            enriched_cves.append(CveEnrichment.enrich_cve(cve_id))
+            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)
+            except Exception as e:
+                raise ValueError(f"{e}")
 
-        return enriched_cves
+        return self
 
+
     splunk_app_enrichment: Optional[List[dict]] = None
 
     @computed_field

contentctl/objects/detection_tags.py

@@ -145,7 +145,7 @@ def serialize_model(self):
     @model_validator(mode="after")
     def addAttackEnrichment(self, info:ValidationInfo):
         if len(self.mitre_attack_enrichments) > 0:
-            raise ValueError(f"Error, field 'mitre_attack_enrichment' should be empty and dynamically populated at runtime. Instead, this field contained: {str(v)}")
+            raise ValueError(f"Error, field 'mitre_attack_enrichment' should be empty and dynamically populated at runtime. Instead, this field contained: {self.mitre_attack_enrichments}")
 
         output_dto:Union[DirectorOutputDto,None]= info.context.get("output_dto",None)
         if output_dto is None: