diff --git a/bin/doc_gen.py b/bin/doc_gen.py index 83c71a4cdb..282436d5b5 100644 --- a/bin/doc_gen.py +++ b/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