Merge pull request #4 from usdAG/fix/error_too_many_sql_variables

Fix issue 'too many SQL variables'
ra1nb0rn · Nov 29, 2023 · eb3f42c · eb3f42c
2 parents 55e7e4c + 26f9fe5
commit eb3f42c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 13 deletions.
diff --git a/cpe_search.py b/cpe_search.py
@@ -2,6 +2,7 @@
 
 import argparse
 from collections import Counter
+from itertools import chain
 import math
 import os
 import pprint
@@ -429,19 +430,25 @@ def words_in_line(words, line):
                 all_cpe_entry_ids.append(eid)
 
     # iterate over all retrieved CPE infos and find best matching CPEs for queries
-    if not all_cpe_entry_ids:
-        iterator = []
-
-    param_in_str = ('?,' * len(all_cpe_entry_ids))[:-1]
-    if keep_data_in_memory:
-        db_query = 'SELECT cpe, term_frequencies, abs_term_frequency FROM cpe_entries WHERE entry_id IN (%s)' % param_in_str
-        cpe_infos = db_cursor.execute(db_query, all_cpe_entry_ids).fetchall()
-        relevant_cpe_infos = cpe_infos
-        iterator = relevant_cpe_infos
-    else:
-        db_query = 'SELECT cpe, term_frequencies, abs_term_frequency FROM cpe_entries WHERE entry_id IN (%s)' % param_in_str
-        db_cursor.execute(db_query, all_cpe_entry_ids)
-        iterator = db_cursor
+    iterator = []
+    max_results_per_query = 250000
+    remaining = len(all_cpe_entry_ids)
+    while remaining > 0:
+        if remaining > max_results_per_query:
+            count_params_in_str = max_results_per_query
+        else:
+            count_params_in_str = remaining
+        param_in_str = ('?,' * count_params_in_str)[:-1]
+        if keep_data_in_memory:
+            db_query = 'SELECT cpe, term_frequencies, abs_term_frequency FROM cpe_entries WHERE entry_id IN (%s)' % param_in_str
+            cpe_infos = db_cursor.execute(db_query, all_cpe_entry_ids[remaining-count_params_in_str:remaining]).fetchall()
+            relevant_cpe_infos = cpe_infos
+            iterator = chain(iterator, relevant_cpe_infos)
+        else:
+            db_query = 'SELECT cpe, term_frequencies, abs_term_frequency FROM cpe_entries WHERE entry_id IN (%s)' % param_in_str
+            db_cursor.execute(db_query, all_cpe_entry_ids[remaining-count_params_in_str:remaining])
+            iterator = chain(iterator, db_cursor)
+        remaining -= max_results_per_query
 
     for cpe_info in iterator:
         cpe, cpe_tf, cpe_abs = cpe_info

diff --git a/test.py b/test.py
@@ -86,5 +86,14 @@ def test_search_electron_1317(self):
         self.assertEqual(result[query][0][0], test_best_match_cpe)
         self.assertEqual(str(result[query][0][1]), test_best_match_score)
 
+    def test_search_blackice_agent_for_server_30(self):
+        self.maxDiff = None
+        query = 'BlackIce Agent for Server 3.0'
+        test_best_match_cpe = 'cpe:2.3:a:iss:blackice_agent_for_server:3.0:*:*:*:*:*:*:*'
+        test_best_match_score = '0.9128709291752767'
+        result = search_cpes(queries=[query])
+        self.assertEqual(result[query][0][0], test_best_match_cpe)
+        self.assertEqual(str(result[query][0][1]), test_best_match_score)
+
 if __name__ == '__main__':
     unittest.main()