From fefcd5138309194b6cf8f42d8cbff1ebb59d7592 Mon Sep 17 00:00:00 2001
From: Palash Oswal <hello@oswalpalash.com>
Date: Wed, 25 Oct 2023 18:12:07 -0400
Subject: [PATCH] Migrate to NVD 2.0 API

---
 securitylist/src/securitylist/NVD.py | 19 ++++++++++---------
 securitylist/src/tests/test_nvd.py   |  4 ++--
 securitylist/src/update_nvd.py       |  2 +-
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/securitylist/src/securitylist/NVD.py b/securitylist/src/securitylist/NVD.py
index 41531f0..f84a0f5 100644
--- a/securitylist/src/securitylist/NVD.py
+++ b/securitylist/src/securitylist/NVD.py
@@ -4,6 +4,7 @@
 import requests
 import datetime
 import time
+import json
 
 class UnexpectedResults(Exception):
     pass
@@ -14,7 +15,7 @@ def __init__(self):
         self.now = datetime.datetime.utcnow()
         self.total = 0
         self.index = 0
-        self.nvd_url = "https://services.nvd.nist.gov/rest/json/cves/1.0"
+        self.nvd_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
         self.payload = {}
         self.last_update = self.now
 
@@ -28,8 +29,8 @@ def __get_time__(self, ts):
         mi = ts.minute
         s = ts.second
 
-        return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}:000 UTC-00:00"
-        #return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}:000+00:00"
+        return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}.000"
+        #return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}.000"
 
     def get_end_time_str(self):
 
@@ -42,12 +43,12 @@ def get_end_time_str(self):
         mi = ts.minute
         s = ts.second
 
-        return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}:000"
+        return f"{y}-{m:02}-{d:02}T{h:02}:{mi:02}:{s:02}.000"
 
     def get_range(self, start, end):
 
         if start is None:
-            self.start_time = datetime.datetime.fromisoformat("1990-01-01T00:00:00:000")
+            self.start_time = datetime.datetime.fromisoformat("1990-01-01T00:00:00.000")
         else:
             self.start_time = datetime.datetime.fromisoformat(start)
 
@@ -82,8 +83,8 @@ def get_page(self, page):
         self.payload = {
             "startIndex": self.index,
             "resultsPerPage": 500,
-            "modStartDate": self.__get_time__(self.start_time),
-            "modEndDate": self.__get_time__(self.end_time)
+            "lastModStartDate": self.__get_time__(self.start_time),
+            "lastModEndDate": self.__get_time__(self.end_time)
         }
 
         response = requests.get(self.nvd_url, params=self.payload)
@@ -100,13 +101,13 @@ def __next__(self):
         if self.iter_n == self.total:
                 raise StopIteration
 
-        if self.iter_current == len(self.data["result"]["CVE_Items"]):
+        if self.iter_current == len(self.data["vulnerabilities"]):
             # Time to paginate
             self.iter_current = 0
             self.page = self.page + 1
             self.get_page(self.page)
 
-        to_return = self.data["result"]["CVE_Items"][self.iter_current]
+        to_return = self.data["vulnerabilities"][self.iter_current]
         self.iter_n = self.iter_n + 1
         self.iter_current = self.iter_current + 1
         return to_return
diff --git a/securitylist/src/tests/test_nvd.py b/securitylist/src/tests/test_nvd.py
index ecb67c7..867d75e 100644
--- a/securitylist/src/tests/test_nvd.py
+++ b/securitylist/src/tests/test_nvd.py
@@ -26,9 +26,9 @@ def json(self):
         def raise_for_status(self):
             pass
 
-    if args[0] == 'https://services.nvd.nist.gov/rest/json/cve/1.0/CVE-1000-0001':
+    if args[0] == 'https://services.nvd.nist.gov/rest/json/cve/2.0/CVE-1000-0001':
         return MockResponse(200)
-    elif args[0] == 'https://services.nvd.nist.gov/rest/json/cve/1.0/CVE-1000-0002':
+    elif args[0] == 'https://services.nvd.nist.gov/rest/json/cve/2.0/CVE-1000-0002':
         return MockResponse(200)
 
     return MockResponse(404)
diff --git a/securitylist/src/update_nvd.py b/securitylist/src/update_nvd.py
index 2687490..dfe76eb 100755
--- a/securitylist/src/update_nvd.py
+++ b/securitylist/src/update_nvd.py
@@ -23,7 +23,7 @@ def main():
     print("Getting %d IDs" % nvd.total)
 
     for i in nvd:
-        the_id = i['cve']['CVE_data_meta']['ID']
+        the_id = i['cve']['id']
         # We need to put these in the NVD namespace
         c = securitylist.CVE(the_id)
         c.add_data('nvd.nist.gov', i)