Replace ujson with orjson

Author: Prabhu Subramanian

requirements.txt

@@ -2,5 +2,5 @@ requests
 appdirs
 tabulate
 msgpack
-ujson
+orjson
 semver

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="appthreat-vulnerability-db",
-    version="1.4.4",
+    version="1.5.0",
     author="Team AppThreat",
     author_email="cloud@appthreat.com",
     description="AppThreat's vulnerability database and package search library with a built-in file based storage. CVE, GitHub, npm are the primary sources of vulnerabilities.",
@@ -14,7 +14,7 @@
     long_description_content_type="text/markdown",
     url="https://github.com/appthreat/vulnerability-db",
     packages=setuptools.find_packages(),
-    install_requires=["requests", "appdirs", "tabulate", "msgpack", "ujson", "semver"],
+    install_requires=["requests", "appdirs", "tabulate", "msgpack", "orjson", "semver"],
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Intended Audience :: Developers",

test/test_db.py

@@ -135,7 +135,7 @@ def test_gha_search_bulk(test_db, test_gha_data):
         if d["details"]["max_affected_version_including"] != "*"
     ]
     res = db.bulk_index_search(tmp_list)
-    assert len(tmp_list)
+    assert len(res)
 
 
 def test_index_search(test_db, test_vuln_data):
@@ -151,9 +151,7 @@ def test_index_search(test_db, test_vuln_data):
     for d in all_data[:40]:
         version = d["details"]["max_affected_version_including"]
         if version and version != "*":
-            tmp_list.append(
-                {"name": d["details"]["package"], "version": version,}
-            )
+            tmp_list.append({"name": d["details"]["package"], "version": version})
     res = db.bulk_index_search(tmp_list)
     assert len(res)
     for r in res:

vdb/lib/gha.py

@@ -12,12 +12,12 @@
 import os
 import re
 
+import orjson
 import requests
-import ujson
 
 from vdb.lib import config as config
 from vdb.lib.nvd import NvdSource
-from vdb.lib.utils import fix_text, get_default_cve_data
+from vdb.lib.utils import get_default_cve_data
 
 logging.basicConfig(
     level=logging.INFO, format="%(levelname)s [%(asctime)s] %(message)s"
@@ -238,7 +238,7 @@ def convert(self, cve_data):
                     cve_id=cve_id,
                     cwe_id="UNKNOWN",
                     assigner=assigner,
-                    references=ujson.dumps(references),
+                    references=orjson.dumps(references).decode("utf-8"),
                     description="",
                     vectorString=vectorString,
                     vendor=vendor.lower(),
@@ -260,10 +260,10 @@ def convert(self, cve_data):
                     lastModifiedDate=cve["updatedAt"],
                 )
                 try:
-                    tdata_json = ujson.loads(tdata)
+                    tdata_json = orjson.loads(tdata)
                     vuln = NvdSource.convert_vuln(tdata_json)
                     vuln.description = description
                     ret_data.append(vuln)
-                except Exception:
-                    LOG.warning(tdata)
+                except Exception as e:
+                    LOG.debug(e)
         return ret_data, page_info

vdb/lib/npm.py

@@ -5,8 +5,8 @@
 """
 import logging
 
+import orjson
 import requests
-import ujson
 
 from vdb.lib import config as config
 from vdb.lib.nvd import NvdSource
@@ -227,7 +227,7 @@ def to_vuln(self, v, ret_data):
                     cve_id=cve_id,
                     cwe_id=cwe_id,
                     assigner=assigner,
-                    references=ujson.dumps(references),
+                    references=orjson.dumps(references).decode("utf-8"),
                     description="",
                     vectorString=vectorString,
                     vendor=vendor,
@@ -249,7 +249,7 @@ def to_vuln(self, v, ret_data):
                     lastModifiedDate=lastModifiedDate,
                 )
                 try:
-                    vuln = NvdSource.convert_vuln(ujson.loads(tdata))
+                    vuln = NvdSource.convert_vuln(orjson.loads(tdata))
                     vuln.description = description
                     ret_data.append(vuln)
                 except Exception:

vdb/lib/nvd.py

@@ -3,8 +3,8 @@
 import logging
 import tempfile
 
+import orjson
 import requests
-import ujson
 
 from vdb.lib import CvssV3, Vulnerability, VulnerabilityDetail, VulnerabilitySource
 from vdb.lib import config as config
@@ -73,7 +73,7 @@ def fetch(self, year):
             with gzip.open(tf.name, "rb") as gzipjf:
                 cve_data = gzipjf.read()
                 try:
-                    json_data = ujson.loads(cve_data)
+                    json_data = orjson.loads(cve_data)
                     return self.convert(json_data)
                 except Exception:
                     logging.warning("Exception while parsing NVD CVE feed")

vdb/lib/utils.py

@@ -1,6 +1,6 @@
 import importlib
 import re
-from datetime import datetime, date
+from datetime import date, datetime
 from enum import Enum
 
 from semver import VersionInfo
@@ -71,7 +71,7 @@ def _load(d):
                 raise ClassNotFoundError(
                     "Class '%s' not found in the given module!" % t
                 )
-            except TypeError as te:
+            except TypeError:
                 raise TypeError(
                     "Make sure there is an constuctor that doesn't take any arguments (class: %s)"
                     % t