From 09c54931af8fd4da5b9d07543f5462766597c211 Mon Sep 17 00:00:00 2001
From: pyth0n1c <87383215+pyth0n1c@users.noreply.github.com>
Date: Fri, 18 Feb 2022 13:36:09 -0800
Subject: [PATCH 1/2] Memoizing/caching a call to CVESSEARCH_API_URL. Gives a
 SIGNIFICANT speedup in docgen.

---
 bin/doc_gen.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/bin/doc_gen.py b/bin/doc_gen.py
index 83c71a4cdb..b34e4bc165 100644
--- a/bin/doc_gen.py
+++ b/bin/doc_gen.py
@@ -1,3 +1,4 @@
+from functools import cache
 import glob
 import yaml
 import argparse
@@ -94,6 +95,13 @@ 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.
+@cache
 def get_cve_enrichment_new(cve_id):
     cve = CVESearch(CVESSEARCH_API_URL)
     result = cve.id(cve_id)

From 46e0bb0a01b16bf1c93b596e238bbc1c8bd23b61 Mon Sep 17 00:00:00 2001
From: pyth0n1c <87383215+pyth0n1c@users.noreply.github.com>
Date: Fri, 18 Feb 2022 14:19:57 -0800
Subject: [PATCH 2/2] Added a timeout to the CVESSEARCH_API enrichment in
 doc_gen.  Sometimes, this was causing doc_gen to freeze.  Now, when the API
 is not working or is extremely slow and a request takes longer than 10
 seconds, a descriptive error will be thrown and the process will exit with 1

---
 bin/doc_gen.py | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/bin/doc_gen.py b/bin/doc_gen.py
index b34e4bc165..282436d5b5 100644
--- a/bin/doc_gen.py
+++ b/bin/doc_gen.py
@@ -1,3 +1,4 @@
+from threading import Thread
 from functools import cache
 import glob
 import yaml
@@ -13,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):
@@ -101,6 +103,9 @@ def parse_and_add_lookups(search_string, lookups):
 #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)
@@ -111,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)
@@ -335,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