Add CVE list word indexing options

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`.
+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.
 5. Take a cup of black or green tea ().
 6. `python3 ./bin/server.py` to run the local HTTP server.
 
@@ -204,14 +204,15 @@ cpe (vendor:product) per version to give a probability of the CPE appearance.
 
 You can also build the ranking directly from the CVE v5 NDJSON dump. The importer scans the full record recursively so it can pick up
 CPE values from both affected-product metadata and vulnerable configuration blocks, then increments the `rank:cpe` and
-`rank:vendor_product` sorted sets using the normalized `cpe:2.3:<part>:<vendor>:<product>` tuple.
+`rank:vendor_product` sorted sets using the normalized `cpe:2.3:<part>:<vendor>:<product>` tuple. It also records any vendor/product words whose `w:<word>` set does not already exist into `set:missing_words_from_cvelistv5`, and `--index-words` lets the CVE v5 importer populate `w:<word>` and `s:<word>` just like the NVD importer.
 
 ### Valkey structure
 
 - `w:<word>` set
 - `s:<word>` sorted set with a score depending of the number of appearance
 - `rank:cpe` sorted set of common `vendor:product` tuples
 - `rank:vendor_product` alias sorted set of the same tuple ranking, populated by both importers
+- `set:missing_words_from_cvelistv5` set of vendor/product words seen in CVE v5 CPEs that were not already present as `w:<word>` keys before the record was processed
 
 ## License
 

bin/import_cvelistv5.py

@@ -22,6 +22,8 @@
 valkey_port = settings.get("valkey.port", 6379)
 valkey_db = settings.get("valkey.db", 8)
 
+DEFAULT_MISSING_WORD_SET = "set:missing_words_from_cvelistv5"
+
 rdb = valkey.Valkey(host=valkey_host, port=valkey_port, db=valkey_db)
 
 
@@ -40,7 +42,27 @@
         "--replace-rank",
         action="store_true",
         default=False,
-        help="Delete existing rank:cpe and rank:vendor_product sorted sets before importing.",
+        help=(
+            "Delete existing rank:cpe, rank:vendor_product, and the missing-word set "
+            "before importing."
+        ),
+    )
+    argparser.add_argument(
+        "--index-words",
+        action="store_true",
+        default=False,
+        help=(
+            "Also index vendor/product words into w:<word> and s:<word> like "
+            "bin/import.py does."
+        ),
+    )
+    argparser.add_argument(
+        "--missing-word-set",
+        default=DEFAULT_MISSING_WORD_SET,
+        help=(
+            "Valkey set name used to store words seen in CVE v5 CPEs that were "
+            "not already present in w:<word>."
+        ),
     )
     args = argparser.parse_args()
 
@@ -60,18 +82,24 @@
         cvelist_file = cvelist_path
 
     if args.replace_rank:
-        removed = rdb.delete("rank:cpe", "rank:vendor_product")
+        removed = rdb.delete("rank:cpe", "rank:vendor_product", args.missing_word_set)
         print(f"Deleted {removed} existing rank key(s) from the database.")
 
-    handler = CVEListV5Handler(rdb)
+    handler = CVEListV5Handler(
+        rdb,
+        index_words=args.index_words,
+        missing_word_key=args.missing_word_set,
+    )
     label = f"{handler.__class__.__name__}[{os.path.basename(cvelist_file)}]"
     print(f"Using {handler.__class__.__name__} to parse file {cvelist_file}...")
     handler.parse_file(cvelist_file, label=label)
 
     rank_size = rdb.zcard("rank:cpe")
     alias_size = rdb.zcard("rank:vendor_product")
+    missing_word_count = rdb.scard(args.missing_word_set)
     print(
         "Done! "
-        f"{rank_size} vendor/product tuples stored in rank:cpe and "
-        f"{alias_size} in rank:vendor_product."
+        f"{rank_size} vendor/product tuples stored in rank:cpe, "
+        f"{alias_size} in rank:vendor_product, and "
+        f"{missing_word_count} words tracked in {args.missing_word_set}."
     )

lib/cpeimport/base.py

@@ -112,6 +112,25 @@ def process_rank_batch(self, cpes, rdb=None):
         pipeline.execute()
         return itemcount, 0
 
+    def collect_missing_words(self, words, missing_word_key, rdb=None):
+        """Record words that are not yet indexed in Valkey."""
+        if not missing_word_key or not words:
+            return 0
+
+        client = rdb or self.rdb
+        pipeline = client.pipeline(transaction=False)
+        new_words = 0
+
+        for word in words:
+            if client.exists(f"w:{word}"):
+                continue
+            pipeline.sadd(missing_word_key, word)
+            new_words += 1
+
+        if new_words:
+            pipeline.execute()
+        return new_words
+
     def record_progress(self, itemcount, wordcount):
         self.itemcount += itemcount
         self.wordcount += wordcount

lib/cpeimport/cvelistv5.py

@@ -15,6 +15,16 @@ class CVEListV5Handler(CPEImportHandler):
         "Expecting property name enclosed in double quotes",
     )
 
+    def __init__(
+        self,
+        rdb,
+        index_words=False,
+        missing_word_key="set:missing_words_from_cvelistv5",
+    ):
+        super().__init__(rdb)
+        self.index_words = index_words
+        self.missing_word_key = missing_word_key
+
     def _parse_impl(self, path):
         if not path.endswith(".ndjson"):
             raise ValueError(f"Unsupported file type: {path}")
@@ -72,7 +82,16 @@ def _process_record(self, record):
             self.skipped += 1
             return
 
-        itemcount, wordcount = self.process_rank_batch(cpes)
+        words = sorted(
+            {word for cpe in cpes for word in self.build_insert_words(cpe)[1]}
+        )
+        self.collect_missing_words(words, self.missing_word_key)
+
+        if self.index_words:
+            itemcount, wordcount = self.process_cpe_batch(cpes)
+        else:
+            itemcount, wordcount = self.process_rank_batch(cpes)
+
         self.record_progress(itemcount=itemcount, wordcount=wordcount)
 
     def _load_record(self, payload):

tests/test_cvelistv5.py

@@ -10,23 +10,34 @@ def __init__(self, rdb):
         self.rdb = rdb
         self.operations = []
 
+    def sadd(self, key, value):
+        self.operations.append(("sadd", key, value))
+
     def zadd(self, key, mapping, incr=False):
         self.operations.append(("zadd", key, mapping, incr))
 
     def execute(self):
         for operation in self.operations:
-            _, key, mapping, incr = operation
-            self.rdb.zadd(key, mapping, incr=incr)
+            if operation[0] == "sadd":
+                _, key, value = operation
+                self.rdb.sadd(key, value)
+            elif operation[0] == "zadd":
+                _, key, mapping, incr = operation
+                self.rdb.zadd(key, mapping, incr=incr)
         self.operations.clear()
 
 
 class FakeRDB:
     def __init__(self):
+        self.sets = {}
         self.sorted_sets = {}
 
     def pipeline(self, transaction=False):
         return FakePipeline(self)
 
+    def sadd(self, key, value):
+        self.sets.setdefault(key, set()).add(value)
+
     def zadd(self, key, mapping, incr=False):
         zset = self.sorted_sets.setdefault(key, {})
         for member, score in mapping.items():
@@ -35,6 +46,9 @@ def zadd(self, key, mapping, incr=False):
             else:
                 zset[member] = score
 
+    def exists(self, key):
+        return key in self.sets or key in self.sorted_sets
+
 
 class CVEListV5HandlerTestCase(unittest.TestCase):
     def test_extracts_cpes_from_metadata_and_configurations(self):
@@ -137,6 +151,67 @@ def test_process_ndjson_counts_unique_vendor_product_tuples_per_record(self):
         self.assertEqual(
             rdb.sorted_sets["rank:vendor_product"]["cpe:2.3:a:acme:widget"], 2
         )
+        self.assertEqual(
+            rdb.sets["set:missing_words_from_cvelistv5"],
+            {"acme", "gadget", "widget"},
+        )
+
+    def test_index_words_adds_w_and_s_entries(self):
+        record = {
+            "containers": {
+                "cna": {
+                    "affected": [
+                        {
+                            "cpes": [
+                                "cpe:2.3:a:acme:rocket_launcher:1.0:*:*:*:*:*:*:*",
+                                "cpe:2.3:a:acme:rocket_launcher:2.0:*:*:*:*:*:*:*",
+                            ]
+                        }
+                    ]
+                }
+            }
+        }
+        payload = io.StringIO(json.dumps(record))
+
+        rdb = FakeRDB()
+        handler = CVEListV5Handler(rdb, index_words=True)
+        handler.process_ndjson_file(payload)
+
+        self.assertEqual(handler.itemcount, 1)
+        self.assertEqual(handler.wordcount, 3)
+        self.assertEqual(rdb.sets["w:acme"], {"cpe:2.3:a:acme:rocket_launcher"})
+        self.assertEqual(rdb.sets["w:rocket"], {"cpe:2.3:a:acme:rocket_launcher"})
+        self.assertEqual(rdb.sets["w:launcher"], {"cpe:2.3:a:acme:rocket_launcher"})
+        self.assertEqual(
+            rdb.sorted_sets["s:rocket"]["cpe:2.3:a:acme:rocket_launcher"], 1
+        )
+        self.assertEqual(
+            rdb.sorted_sets["rank:cpe"]["cpe:2.3:a:acme:rocket_launcher"], 3
+        )
+
+    def test_only_tracks_missing_words_absent_from_existing_index(self):
+        record = {
+            "containers": {
+                "cna": {
+                    "affected": [
+                        {
+                            "cpes": [
+                                "cpe:2.3:a:acme:widget:1.0:*:*:*:*:*:*:*",
+                                "cpe:2.3:a:acme:gadget:1.0:*:*:*:*:*:*:*",
+                            ]
+                        }
+                    ]
+                }
+            }
+        }
+
+        rdb = FakeRDB()
+        rdb.sadd("w:acme", "cpe:2.3:a:acme:existing")
+        rdb.sadd("w:widget", "cpe:2.3:a:acme:existing")
+        handler = CVEListV5Handler(rdb)
+        handler.process_ndjson_file(io.StringIO(json.dumps(record)))
+
+        self.assertEqual(rdb.sets["set:missing_words_from_cvelistv5"], {"gadget"})
 
     def test_skips_single_invalid_multiline_record_and_continues(self):
         record_one = {