Branch was auto-updated.

bin/doc_gen.py

@@ -1,3 +1,5 @@
+from threading import Thread
+from functools import cache
 import glob
 import yaml
 import argparse
@@ -12,6 +14,7 @@ from stix2 import Filter
 from pycvesearch import CVESearch
 
 CVESSEARCH_API_URL = 'https://cve.circl.lu'
+CVESSEARCH_API_TIMEOUT = 10
 
 
 def load_objects(REPO_PATH, TYPE):
@@ -94,6 +97,16 @@ def parse_and_add_lookups(search_string, lookups):
     return lookup_objects
 
 
+#This function hits an API, which can be slow, especially if the API
+#or network connection is slow. This has sometimes caused issues
+#in our CI/CD taking a very long time.  Here, we memoize/cache
+#the calls to this function - there will be many duplicates which
+#should all return the same results. This will gives us a significant
+#speedup.
+#Hackish way of keeping the original functionality and allowing it
+#to easily return a value when in a thread... making an optional
+#default argument with a type that is mutable!
+@cache
 def get_cve_enrichment_new(cve_id):
     cve = CVESearch(CVESSEARCH_API_URL)
     result = cve.id(cve_id)
@@ -103,6 +116,15 @@ def get_cve_enrichment_new(cve_id):
     cve_enriched['summary'] = result['summary']
     return cve_enriched
 
+#helper function to easily return a value from a thread that is 
+#running a memoized/@cached function
+def get_cve_enrichment_new_wrapper(cve_id,mutable_list):
+    mutable_list.append(get_cve_enrichment_new(cve_id))
+    
+
+
+
+
 def get_all_techniques(projects_path):
     path_cti = path.join(projects_path,'cti/enterprise-attack')
     fs = FileSystemSource(path_cti)
@@ -327,8 +349,17 @@ def generate_doc_detections(REPO_PATH, OUTPUT_DIR, TEMPLATE_PATH, attack, messag
         cves = []
         if 'cve' in detection_yaml['tags']:
             for cve_id in detection_yaml['tags']['cve']:
-                cve = get_cve_enrichment_new(cve_id)
-                cves.append(cve)
+                mutable_list = []
+                try:
+                    cve_thread = Thread(target=get_cve_enrichment_new_wrapper, args=(cve_id, mutable_list))
+                    cve_thread.start()
+                    cve_thread.join(timeout=CVESSEARCH_API_TIMEOUT)
+                    if cve_thread.is_alive():
+                        raise(Exception(f"Timed out getting CVE Enrichment from {CVESSEARCH_API_URL} after {CVESSEARCH_API_TIMEOUT} seconds."))
+                    cves.append(mutable_list[0])
+                except Exception as e:
+                    print(f"Error - {str(e)}\nQuitting...",file=sys.stderr)
+                    sys.exit(1)
             detection_yaml['cve'] = cves
 
         # enrich with macros