Reset CVE v5 rank state before import

Author: adulau

README.md

@@ -24,7 +24,7 @@ curl -s -X POST http://localhost:8000/search -d '{"query": ["tomcat"]}' | jq .
 1. `git clone https://github.com/cve-search/cpe-guesser.git`
 2. `cd cpe-guesser`
 3. Download the CPE dictionary & populate the database with `python3 ./bin/import.py` (the NVD JSON importer now uses parallel workers by default; tune with `--workers` and `--batch-size` if needed).
-4. Optionally enrich the vendor/product ranking with CVE v5 data using `python3 ./bin/import_cvelistv5.py`. This downloads `./data/cvelistv5.ndjson` only when missing unless you pass `--download`. Use `--index-words` if you also want to populate `w:<word>` and `s:<word>` from CVE v5 data, and review `set:missing_words_from_cvelistv5` for words that were not already indexed in Valkey.
+4. Optionally enrich the vendor/product ranking with CVE v5 data using `python3 ./bin/import_cvelistv5.py`. This downloads `./data/cvelistv5.ndjson` only when missing unless you pass `--download`, and it now resets `rank:cpe`, `rank:vendor_product`, and `set:missing_words_from_cvelistv5` before each import unless you opt into `--preserve-rank`. Use `--index-words` if you also want to populate `w:<word>` and `s:<word>` from CVE v5 data, and review `set:missing_words_from_cvelistv5` for words that were not already indexed in Valkey.
 5. Take a cup of black or green tea ().
 6. `python3 ./bin/server.py` to run the local HTTP server.
 

bin/import_cvelistv5.py

@@ -10,7 +10,7 @@
 from dynaconf import Dynaconf
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
-from lib.cpeimport import CPEDownloader, CVEListV5Handler
+from lib.cpeimport import CPEDownloader, CVEListV5Handler, reset_rank_state
 
 # Configuration
 settings = Dynaconf(settings_files=["../config/settings.yaml"])
@@ -39,12 +39,12 @@
         help="Force downloading the CVE v5 NDJSON even if it already exists locally.",
     )
     argparser.add_argument(
-        "--replace-rank",
+        "--preserve-rank",
         action="store_true",
         default=False,
         help=(
-            "Delete existing rank:cpe, rank:vendor_product, and the missing-word set "
-            "before importing."
+            "Keep existing rank:cpe, rank:vendor_product, and the missing-word set "
+            "instead of resetting them before importing."
         ),
     )
     argparser.add_argument(
@@ -81,9 +81,14 @@
         print(f"Using existing file {cvelist_path} ...")
         cvelist_file = cvelist_path
 
-    if args.replace_rank:
-        removed = rdb.delete("rank:cpe", "rank:vendor_product", args.missing_word_set)
-        print(f"Deleted {removed} existing rank key(s) from the database.")
+    if args.preserve_rank:
+        print(
+            "Preserving existing rank:cpe, rank:vendor_product, and "
+            f"{args.missing_word_set}."
+        )
+    else:
+        removed = reset_rank_state(rdb, args.missing_word_set)
+        print(f"Reset {removed} existing rank key(s) before importing.")
 
     handler = CVEListV5Handler(
         rdb,

lib/cpeimport/__init__.py

@@ -1,12 +1,13 @@
 from .base import CPEImportHandler
-from .cvelistv5 import CVEListV5Handler
+from .cvelistv5 import CVEListV5Handler, reset_rank_state
 from .nvd_json import NVDCPEHandler
 from .xml_dictionary import XMLCPEHandler
 from .downloader import CPEDownloader
 
 __all__ = [
     "CPEImportHandler",
     "CVEListV5Handler",
+    "reset_rank_state",
     "NVDCPEHandler",
     "XMLCPEHandler",
     "CPEDownloader",

lib/cpeimport/cvelistv5.py

@@ -3,6 +3,11 @@
 from .base import CPEImportHandler
 
 
+def reset_rank_state(rdb, missing_word_key="set:missing_words_from_cvelistv5"):
+    """Delete the CVE v5 ranking keys so each import starts from a clean state."""
+    return rdb.delete("rank:cpe", "rank:vendor_product", missing_word_key)
+
+
 class CVEListV5Handler(CPEImportHandler):
     """Handler for CVE Record Format v5 NDJSON exports."""
 

tests/test_cvelistv5.py

@@ -2,7 +2,7 @@
 import json
 import unittest
 
-from lib.cpeimport.cvelistv5 import CVEListV5Handler
+from lib.cpeimport.cvelistv5 import CVEListV5Handler, reset_rank_state
 
 
 class FakePipeline:
@@ -32,6 +32,14 @@ def __init__(self):
         self.sets = {}
         self.sorted_sets = {}
 
+    def delete(self, *keys):
+        removed = 0
+        for key in keys:
+            removed += int(key in self.sets or key in self.sorted_sets)
+            self.sets.pop(key, None)
+            self.sorted_sets.pop(key, None)
+        return removed
+
     def pipeline(self, transaction=False):
         return FakePipeline(self)
 
@@ -51,6 +59,21 @@ def exists(self, key):
 
 
 class CVEListV5HandlerTestCase(unittest.TestCase):
+    def test_reset_rank_state_clears_previous_import_data(self):
+        rdb = FakeRDB()
+        rdb.zadd("rank:cpe", {"cpe:2.3:a:acme:widget": 3})
+        rdb.zadd("rank:vendor_product", {"cpe:2.3:a:acme:widget": 3})
+        rdb.sadd("set:missing_words_from_cvelistv5", "widget")
+        rdb.sadd("w:widget", "cpe:2.3:a:acme:widget")
+
+        removed = reset_rank_state(rdb)
+
+        self.assertEqual(removed, 3)
+        self.assertNotIn("rank:cpe", rdb.sorted_sets)
+        self.assertNotIn("rank:vendor_product", rdb.sorted_sets)
+        self.assertNotIn("set:missing_words_from_cvelistv5", rdb.sets)
+        self.assertIn("w:widget", rdb.sets)
+
     def test_extracts_cpes_from_metadata_and_configurations(self):
         record = {
             "containers": {