Refactor CPE imports and add NVDCPEHandler for JSON feeds

.gitignore

@@ -1,4 +1,7 @@
 bin/__pycache__/
 lib/__pycache__/
-data/*.xml
+lib/cpeimport/__pycache__/
 data/*.gz
+data/*.tar
+data/*.json
+data/*.xml

README.md

@@ -153,6 +153,34 @@ python3 lookup.py microsoft sql server | jq .
 
 A CPE entry is composed of a human readable name with some references and the structured CPE name.
 
+NVD CPE Dictionary 2.0 JSON feeds:
+
+```json
+    {
+      "cpe": {
+        "deprecated": false,
+        "cpeName": "cpe:2.3:a:10web:form_maker:1.7.17:*:*:*:*:wordpress:*:*",
+        "cpeNameId": "A6F0EC33-CD55-4FE2-8B55-5C0F9CC88BF9",
+        "lastModified": "2019-06-04T16:47:17.300",
+        "created": "2019-06-04T16:47:17.300",
+        "titles": [
+          {
+            "title": "10web Form Maker 1.7.17 for WordPress",
+            "lang": "en"
+          }
+        ],
+        "refs": [
+          {
+            "ref": "https://wordpress.org/plugins/form-maker/#developers",
+            "type": "Change Log"
+          }
+        ]
+      }
+    },
+```
+
+Legacy XML feeds:
+
 ```xml
   <cpe-item name="cpe:/a:10web:form_maker:1.7.17::~~~wordpress~~">
     <title xml:lang="en-US">10web Form Maker 1.7.17 for WordPress</title>
@@ -183,8 +211,8 @@ cpe (vendor:product) per version to give a probability of the CPE appearance.
 Software is open source and released under a 2-Clause BSD License
 
 ~~~
-  Copyright (C) 2021-2024 Alexandre Dulaunoy
-  Copyright (C) 2021-2024 Esa Jokinen
+  Copyright (C) 2021-2025 Alexandre Dulaunoy
+  Copyright (C) 2021-2025 Esa Jokinen
 ~~~
 
 We welcome contributions! All contributors collectively own the CPE Guesser project. By contributing, contributors also acknowledge the [Developer Certificate of Origin](https://developercertificate.org/) when submitting pull requests or using other methods of contribution.

bin/import.py

@@ -4,20 +4,19 @@
 import argparse
 import sys
 import os
-import urllib.request
-import gzip
-import shutil
-import xml.sax
 import valkey
-import time
+import urllib.error
 from dynaconf import Dynaconf
 
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from lib.cpeimport import CPEDownloader, NVDCPEHandler, XMLCPEHandler
+
 # Configuration
 settings = Dynaconf(settings_files=["../config/settings.yaml"])
-cpe_path = settings.get("cpe.path", "../data/official-cpe-dictionary_v2.3.xml")
+cpe_path = settings.get("cpe.path", "./data/nvdcpe-2.0.tar")
 cpe_source = settings.get(
     "cpe.source",
-    "https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz",
+    "https://nvd.nist.gov/feeds/json/cpe/2.0/nvdcpe-2.0.tar.gz",
 )
 valkey_host = settings.get("valkey.host", "127.0.0.1")
 valkey_port = settings.get("valkey.port", 6379)
@@ -26,149 +25,65 @@
 rdb = valkey.Valkey(host=valkey_host, port=valkey_port, db=valkey_db)
 
 
-class CPEHandler(xml.sax.ContentHandler):
-    def __init__(self):
-        self.cpe = ""
-        self.title = ""
-        self.title_seen = False
-        self.cpe = ""
-        self.record = {}
-        self.refs = []
-        self.itemcount = 0
-        self.wordcount = 0
-        self.start_time = time.time()
-
-    def startElement(self, tag, attributes):
-        self.CurrentData = tag
-        if tag == "cpe-23:cpe23-item":
-            self.record["cpe-23"] = attributes["name"]
-        if tag == "title":
-            self.title_seen = True
-        if tag == "reference":
-            self.refs.append(attributes["href"])
-
-    def characters(self, data):
-        if self.title_seen:
-            self.title = self.title + data
-
-    def endElement(self, tag):
-        if tag == "title":
-            self.record["title"] = self.title
-            self.title = ""
-            self.title_seen = False
-        if tag == "references":
-            self.record["refs"] = self.refs
-            self.refs = []
-        if tag == "cpe-item":
-            to_insert = CPEExtractor(cpe=self.record["cpe-23"])
-            for word in canonize(to_insert["vendor"]):
-                insert(word=word, cpe=to_insert["cpeline"])
-                self.wordcount += 1
-            for word in canonize(to_insert["product"]):
-                insert(word=word, cpe=to_insert["cpeline"])
-                self.wordcount += 1
-            self.record = {}
-            self.itemcount += 1
-            if self.itemcount % 5000 == 0:
-                time_elapsed = round(time.time() - self.start_time)
-                print(
-                    f"... {self.itemcount} items processed ({self.wordcount} words) in {time_elapsed} seconds"
-                )
-
-
-def CPEExtractor(cpe=None):
-    if cpe is None:
-        return False
-    record = {}
-    cpefield = cpe.split(":")
-    record["vendor"] = cpefield[3]
-    record["product"] = cpefield[4]
-    cpeline = ""
-    for cpeentry in cpefield[:5]:
-        cpeline = f"{cpeline}:{cpeentry}"
-    record["cpeline"] = cpeline[1:]
-    return record
-
-
-def canonize(value=None):
-    value = value.lower()
-    words = value.split("_")
-    return words
-
-
-def insert(word=None, cpe=None):
-    if cpe is None or word is None:
-        return False
-    rdb.sadd(f"w:{word}", cpe)
-    rdb.zadd(f"s:{word}", {cpe: 1}, incr=True)
-    rdb.zadd("rank:cpe", {cpe: 1}, incr=True)
-
-
 if __name__ == "__main__":
     argparser = argparse.ArgumentParser(
         description="Initializes the Redis database with CPE dictionary."
     )
     argparser.add_argument(
         "--download",
         "-d",
-        action="count",
-        default=0,
+        action="store_true",
+        default=False,
         help="Download the CPE dictionary even if it already exists.",
     )
     argparser.add_argument(
         "--replace",
         "-r",
-        action="count",
-        default=0,
-        help="Flush and repopulated the CPE database.",
-    )
-    argparser.add_argument(
-        "--update",
-        "-u",
         action="store_true",
         default=False,
-        help="Update the CPE database without flushing",
+        help="Flush and repopulated the CPE database.",
     )
     args = argparser.parse_args()
 
-    if args.replace == 0 and rdb.dbsize() > 0 and not args.update:
+    if not args.replace and rdb.dbsize() > 0:
         print(f"Warning! The Redis database already has {rdb.dbsize()} keys.")
         print("Use --replace if you want to flush the database and repopulate it.")
         sys.exit(0)
 
-    if args.download > 0 or not os.path.isfile(cpe_path):
-        print(f"Downloading CPE data from {cpe_source} ...")
+    if args.download or not os.path.isfile(cpe_path):
+        downloader = CPEDownloader(url=cpe_source, dest_path=cpe_path)
         try:
-            urllib.request.urlretrieve(cpe_source, f"{cpe_path}.gz")
+            cpe_file = downloader.download(force=args.download)
         except (
             urllib.error.HTTPError,
             urllib.error.URLError,
             FileNotFoundError,
             PermissionError,
         ) as e:
-            print(e)
-            sys.exit(1)
-
-        print(f"Uncompressing {cpe_path}.gz ...")
-        try:
-            with gzip.open(f"{cpe_path}.gz", "rb") as cpe_gz:
-                with open(cpe_path, "wb") as cpe_xml:
-                    shutil.copyfileobj(cpe_gz, cpe_xml)
-            os.remove(f"{cpe_path}.gz")
-        except (FileNotFoundError, PermissionError) as e:
-            print(e)
             sys.exit(1)
 
     elif os.path.isfile(cpe_path):
         print(f"Using existing file {cpe_path} ...")
+        cpe_file = cpe_path
 
-    if rdb.dbsize() > 0 and not args.update:
+    if rdb.dbsize() > 0 and args.replace:
         print(f"Flushing {rdb.dbsize()} keys from the database...")
         rdb.flushdb()
 
     print("Populating the database (please be patient)...")
-    parser = xml.sax.make_parser()
-    Handler = CPEHandler()
-    parser.setContentHandler(Handler)
-    parser.parse(cpe_path)
+
+    _, ext = os.path.splitext(cpe_file)
+    ext = ext.lower()
+    if ext == ".tar" or ext == ".json":
+        handler = NVDCPEHandler(rdb)
+    elif ext == ".xml":
+        handler = XMLCPEHandler(rdb)
+    else:
+        print(f"Error! No handler for the file type of {cpe_file}")
+        sys.exit(1)
+
+    print(f"Using {handler.__class__.__name__} to parse file {cpe_file}...")
+    label = f"{handler.__class__.__name__}[{os.path.basename(cpe_file)}]"
+    handler.parse_file(cpe_file, label=label)
+
     print(f"Done! {rdb.dbsize()} keys inserted.")

config/settings.yaml

@@ -5,5 +5,5 @@ valkey:
   port: 6379
   db: 8
 cpe:
-  path: '../data/official-cpe-dictionary_v2.3.xml'
-  source: 'https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz'
+  path: './data/nvdcpe-2.0.tar'
+  source: 'https://nvd.nist.gov/feeds/json/cpe/2.0/nvdcpe-2.0.tar.gz'

docker/settings.yaml

@@ -5,5 +5,5 @@ valkey:
   port: 6379
   db: 8
 cpe:
-  path: '/data/official-cpe-dictionary_v2.3.xml'
-  source: 'https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz'
+  path: '/data/nvdcpe-2.0.tar'
+  source: 'https://nvd.nist.gov/feeds/json/cpe/2.0/nvdcpe-2.0.tar.gz'

lib/cpeimport/__init__.py

@@ -0,0 +1,11 @@
+from .base import CPEImportHandler
+from .nvd_json import NVDCPEHandler
+from .xml_dictionary import XMLCPEHandler
+from .downloader import CPEDownloader
+
+__all__ = [
+    "CPEImportHandler",
+    "NVDCPEHandler",
+    "XMLCPEHandler",
+    "CPEDownloader",
+]

lib/cpeimport/base.py

@@ -0,0 +1,70 @@
+import time
+from abc import ABC, abstractmethod
+
+
+class CPEImportHandler(ABC):
+    """Base class with common functionality for importing CPE data."""
+
+    def __init__(self, rdb):
+        self.rdb = rdb
+        self.itemcount = 0
+        self.wordcount = 0
+        self.skipped = 0
+        self.start_time = time.time()
+
+    @abstractmethod
+    def _parse_impl(self, filepath):
+        """Subclasses implement parsing logic for their format."""
+        pass
+
+    def parse_file(self, filepath, label=""):
+        """Common entry point for all handlers."""
+        self.start_time = time.time()
+        self.itemcount = 0
+        self.wordcount = 0
+        self.skipped = 0
+
+        self._parse_impl(filepath)
+
+        elapsed = round(time.time() - self.start_time)
+        msg = f"Finished {label}: {self.itemcount} items " f"({self.wordcount} words)"
+        if self.skipped:
+            msg += f", {self.skipped} skipped"
+        msg += f" in {elapsed} seconds."
+        print(msg)
+
+    def CPEExtractor(self, cpe):
+        fields = cpe.split(":")
+        vendor = fields[3]
+        product = fields[4]
+        cpeline = ":".join(fields[:5])
+        return {"vendor": vendor, "product": product, "cpeline": cpeline}
+
+    def canonize(self, value):
+        return value.lower().split("_")
+
+    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)
+
+    def process_cpe(self, cpe):
+        """Shared vendor/product → Redis word indexing logic."""
+        to_insert = self.CPEExtractor(cpe=cpe)
+
+        for word in self.canonize(to_insert["vendor"]):
+            self.insert(word=word, cpe=to_insert["cpeline"])
+            self.wordcount += 1
+
+        for word in self.canonize(to_insert["product"]):
+            self.insert(word=word, cpe=to_insert["cpeline"])
+            self.wordcount += 1
+
+        self.itemcount += 1
+
+        if self.itemcount % 5000 == 0:
+            time_elapsed = round(time.time() - self.start_time)
+            print(
+                f"... {self.itemcount} items processed "
+                f"({self.wordcount} words) in {time_elapsed} seconds"
+            )