Update cpe_search with further speedup

Also removes memory-based CPE search and increases number of retrieved related queries to 5.
usdAG · Oct 11, 2023 · 904a3ec · 904a3ec
1 parent 061893d
commit 904a3ec
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/cpe_search b/cpe_search
diff --git a/search_vulns.py b/search_vulns.py
@@ -345,16 +345,18 @@ def search_vulns(query, db_cursor=None, software_match_threshold=CPE_SEARCH_THRE
     close_cursor_after = False
     if not db_cursor:
         db_conn_file = sqlite3.connect(DATABASE_FILE)
-        db_conn_mem = sqlite3.connect(':memory:')
-        db_conn_file.backup(db_conn_mem)
-        db_cursor = db_conn_mem.cursor()
-        # db_cursor = db_conn_file.cursor()
+        if keep_data_in_memory:
+            db_conn_mem = sqlite3.connect(':memory:')
+            db_conn_file.backup(db_conn_mem)
+            db_cursor = db_conn_mem.cursor()
+        else:
+            db_cursor = db_conn_file.cursor()
         close_cursor_after = True
 
     # if given query is not already a CPE, retrieve a CPE that matches the query
     cpe = query
     if not MATCH_CPE_23_RE.match(query):
-        cpe = search_cpes(query, count=1, threshold=software_match_threshold, keep_data_in_memory=keep_data_in_memory)
+        cpe = search_cpes(query, count=1, threshold=software_match_threshold)
 
         if not cpe or not cpe[query]:
             return None
@@ -365,7 +367,7 @@ def search_vulns(query, db_cursor=None, software_match_threshold=CPE_SEARCH_THRE
 
         cpe = cpe[query][0][0]
     elif not is_good_cpe:
-        pot_matching_cpe = match_cpe23_to_cpe23_from_dict(cpe, keep_data_in_memory=keep_data_in_memory)
+        pot_matching_cpe = match_cpe23_to_cpe23_from_dict(cpe)
         if pot_matching_cpe:
             cpe = pot_matching_cpe
 
@@ -394,7 +396,7 @@ def search_vulns_return_cpe(query, db_cursor=None, software_match_threshold=CPE_
     cpe, pot_cpes = query, []
     if not MATCH_CPE_23_RE.match(query):
         is_good_cpe = False
-        cpes = search_cpes(query, count=4, threshold=0.25, zero_extend_versions=zero_extend_versions, keep_data_in_memory=keep_data_in_memory)
+        cpes = search_cpes(query, count=5, threshold=0.25, zero_extend_versions=zero_extend_versions)
 
         if not cpes or not cpes[query]:
             return {query: {'cpe': None, 'vulns': None, 'pot_cpes': []}}
@@ -512,9 +514,6 @@ def main():
 
     # get handle for vulnerability database
     db_conn_file = sqlite3.connect(DATABASE_FILE)
-    # db_conn_mem = sqlite3.connect(':memory:')
-    # db_conn_file.backup(db_conn_mem)
-    # db_cursor = db_conn_mem.cursor()
     db_cursor = db_conn_file.cursor()
 
     # retrieve known vulnerabilities for every query and print them