Merge pull request #24 from vulnerability-lookup/codex/add-cpe-extraction-component-in-cpe-guesser

Add CVE v5 NDJSON rank importer and CVEListV5Handler

Author: adulau

README.md

@@ -24,8 +24,9 @@ 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. Take a cup of black or green tea ().
-5. `python3 ./bin/server.py` to run the local HTTP server.
+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`.
+5. Take a cup of black or green tea ().
+6. `python3 ./bin/server.py` to run the local HTTP server.
 
 If you don't want to install it locally, there is a public online version. Check below.
 
@@ -201,10 +202,16 @@ Split vendor name and product name (such as `_`) into single word(s) and then ca
 the cpe vendor:product format as value and the canonized word as key.  Then cpe guesser creates a ranked set with the most common
 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.
+
 ### 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
 
 ## License
 

bin/import_cvelistv5.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import argparse
+import os
+import sys
+import urllib.error
+
+import valkey
+from dynaconf import Dynaconf
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from lib.cpeimport import CPEDownloader, CVEListV5Handler
+
+# Configuration
+settings = Dynaconf(settings_files=["../config/settings.yaml"])
+cvelist_path = settings.get("cvelistv5.path", "./data/cvelistv5.ndjson")
+cvelist_source = settings.get(
+    "cvelistv5.source", "https://vulnerability.circl.lu/dumps/cvelistv5.ndjson"
+)
+valkey_host = settings.get("valkey.host", "127.0.0.1")
+valkey_port = settings.get("valkey.port", 6379)
+valkey_db = settings.get("valkey.db", 8)
+
+rdb = valkey.Valkey(host=valkey_host, port=valkey_port, db=valkey_db)
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser(
+        description="Populate Valkey vendor/product tuple rankings from CVE v5 NDJSON."
+    )
+    argparser.add_argument(
+        "--download",
+        "-d",
+        action="store_true",
+        default=False,
+        help="Force downloading the CVE v5 NDJSON even if it already exists locally.",
+    )
+    argparser.add_argument(
+        "--replace-rank",
+        action="store_true",
+        default=False,
+        help="Delete existing rank:cpe and rank:vendor_product sorted sets before importing.",
+    )
+    args = argparser.parse_args()
+
+    if args.download or not os.path.isfile(cvelist_path):
+        downloader = CPEDownloader(url=cvelist_source, dest_path=cvelist_path)
+        try:
+            cvelist_file = downloader.download(force=args.download)
+        except (
+            urllib.error.HTTPError,
+            urllib.error.URLError,
+            FileNotFoundError,
+            PermissionError,
+        ):
+            sys.exit(1)
+    else:
+        print(f"Using existing file {cvelist_path} ...")
+        cvelist_file = cvelist_path
+
+    if args.replace_rank:
+        removed = rdb.delete("rank:cpe", "rank:vendor_product")
+        print(f"Deleted {removed} existing rank key(s) from the database.")
+
+    handler = CVEListV5Handler(rdb)
+    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")
+    print(
+        "Done! "
+        f"{rank_size} vendor/product tuples stored in rank:cpe and "
+        f"{alias_size} in rank:vendor_product."
+    )

config/settings.yaml

@@ -7,3 +7,6 @@ valkey:
 cpe:
   path: './data/nvdcpe-2.0.tar'
   source: 'https://nvd.nist.gov/feeds/json/cpe/2.0/nvdcpe-2.0.tar.gz'
+cvelistv5:
+  path: './data/cvelistv5.ndjson'
+  source: 'https://vulnerability.circl.lu/dumps/cvelistv5.ndjson'

lib/cpeimport/__init__.py

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

lib/cpeimport/base.py

@@ -37,9 +37,22 @@ def parse_file(self, filepath, label=""):
 
     def CPEExtractor(self, cpe):
         fields = cpe.split(":")
-        vendor = fields[3]
-        product = fields[4]
-        cpeline = ":".join(fields[:5])
+        if cpe.startswith("cpe:/"):
+            if len(fields) < 4:
+                raise ValueError(f"Invalid legacy CPE: {cpe}")
+            part = fields[1].lstrip("/") or "*"
+            vendor = fields[2]
+            product = fields[3]
+            cpeline = f"cpe:2.3:{part}:{vendor}:{product}"
+        else:
+            if len(fields) < 5:
+                raise ValueError(f"Invalid CPE 2.3 entry: {cpe}")
+            vendor = fields[3]
+            product = fields[4]
+            cpeline = ":".join(fields[:5])
+
+        if not vendor or not product:
+            raise ValueError(f"Invalid vendor/product tuple in CPE: {cpe}")
         return {"vendor": vendor, "product": product, "cpeline": cpeline}
 
     def canonize(self, value):
@@ -49,6 +62,7 @@ def insert(self, word, cpe):
         self.rdb.sadd(f"w:{word}", cpe)
         self.rdb.zadd(f"s:{word}", {cpe: 1}, incr=True)
         self.rdb.zadd("rank:cpe", {cpe: 1}, incr=True)
+        self.rdb.zadd("rank:vendor_product", {cpe: 1}, incr=True)
 
     def build_insert_words(self, cpe):
         to_insert = self.CPEExtractor(cpe=cpe)
@@ -73,12 +87,31 @@ def process_cpe_batch(self, cpes, rdb=None):
                 pipeline.sadd(f"w:{word}", cpeline)
                 pipeline.zadd(f"s:{word}", {cpeline: 1}, incr=True)
                 pipeline.zadd("rank:cpe", {cpeline: 1}, incr=True)
+                pipeline.zadd("rank:vendor_product", {cpeline: 1}, incr=True)
                 wordcount += 1
             itemcount += 1
 
         pipeline.execute()
         return itemcount, wordcount
 
+    def process_rank_batch(self, cpes, rdb=None):
+        """Insert only vendor/product tuple ranking data for a batch of CPEs."""
+        if not cpes:
+            return 0, 0
+
+        client = rdb or self.rdb
+        pipeline = client.pipeline(transaction=False)
+        itemcount = 0
+
+        for cpe in cpes:
+            cpeline = self.CPEExtractor(cpe=cpe)["cpeline"]
+            pipeline.zadd("rank:cpe", {cpeline: 1}, incr=True)
+            pipeline.zadd("rank:vendor_product", {cpeline: 1}, incr=True)
+            itemcount += 1
+
+        pipeline.execute()
+        return itemcount, 0
+
     def record_progress(self, itemcount, wordcount):
         self.itemcount += itemcount
         self.wordcount += wordcount

lib/cpeimport/cvelistv5.py

@@ -0,0 +1,64 @@
+import json
+
+from .base import CPEImportHandler
+
+
+class CVEListV5Handler(CPEImportHandler):
+    """Handler for CVE Record Format v5 NDJSON exports."""
+
+    CPE_PREFIXES = ("cpe:2.3:", "cpe:/")
+
+    def _parse_impl(self, path):
+        if not path.endswith(".ndjson"):
+            raise ValueError(f"Unsupported file type: {path}")
+
+        with open(path, "r", encoding="utf-8") as f:
+            self.process_ndjson_file(f)
+
+    def process_ndjson_file(self, fileobj):
+        for line_number, line in enumerate(fileobj, start=1):
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as e:
+                print(f"Skipping invalid NDJSON record on line {line_number}: {e}")
+                self.skipped += 1
+                continue
+
+            cpes = self.extract_cpes(record)
+            if not cpes:
+                self.skipped += 1
+                continue
+
+            itemcount, wordcount = self.process_rank_batch(cpes)
+            self.record_progress(itemcount=itemcount, wordcount=wordcount)
+
+    def extract_cpes(self, record):
+        cpes = set()
+        self._collect_cpes(record, cpes)
+        valid = []
+
+        for cpe in cpes:
+            try:
+                valid.append(self.CPEExtractor(cpe)["cpeline"])
+            except (IndexError, ValueError):
+                self.skipped += 1
+
+        return sorted(set(valid))
+
+    def _collect_cpes(self, node, cpes):
+        if isinstance(node, dict):
+            for value in node.values():
+                self._collect_cpes(value, cpes)
+            return
+
+        if isinstance(node, list):
+            for item in node:
+                self._collect_cpes(item, cpes)
+            return
+
+        if isinstance(node, str) and node.startswith(self.CPE_PREFIXES):
+            cpes.add(node)

lib/cpeimport/downloader.py

@@ -18,6 +18,10 @@ def download(self, force=False):
             print(f"Using existing file {self.dest_path} ...")
             return self.dest_path
 
+        dest_dir = os.path.dirname(self.dest_path)
+        if dest_dir:
+            os.makedirs(dest_dir, exist_ok=True)
+
         print(f"Downloading CPE data from {self.url} ...")
         download_path = (
             self.dest_path + ".gz" if self.url.endswith(".gz") else self.dest_path