get-nvd-json.py

#!/usr/bin/env python

import requests
from requests.adapters import HTTPAdapter
from urllib3.util import Retry
import os
import sys
import json
from tqdm import tqdm

def main():

    # Define the retry strategy
    retry_strategy = Retry(
        total=6,  # Maximum number of retries
        backoff_factor=4,
        status_forcelist=[429, 500, 502, 503, 504],  # HTTP status codes to retry on
    )
    # Create an HTTP adapter with the retry strategy and mount it to session
    adapter = HTTPAdapter(max_retries=retry_strategy)

    # Create a new session object
    session = requests.Session()
    session.mount('http://', adapter)
    session.mount('https://', adapter)

    #nvd_url = "https://services.nvd.nist.gov/rest/json/cves/1.0"
    nvd_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
    payload = {
        "startIndex": 0,
        "resultsPerPage": 2000
    }

    headers = {}

    if "NVD_API_KEY" in os.environ:

        headers['apiKey'] = os.environ['NVD_API_KEY']

    with open("nvd-out.json", "w") as fh:

        # Get page 1
        req = session.get(nvd_url, params=payload, headers=headers)
        if req.status_code != 200:
            print(req)
            sys.exit(1)
        response = req.json()
        results = response["vulnerabilities"]

        for i in results:
            fh.write(json.dumps(i["cve"]))
            fh.write("\n")

        total = response["totalResults"]

        payload["startIndex"] = payload["startIndex"] + response["resultsPerPage"]

        progress_bar = tqdm(total=total)
        progress_bar.update(response["resultsPerPage"])

        #Loop over pages
        while response["resultsPerPage"] > 0:
            req = session.get(nvd_url, params=payload, headers=headers)
            if req.status_code != 200:
                print(req)
                sys.exit(1)
            response = req.json()
            payload["startIndex"] = payload["startIndex"] + response["resultsPerPage"]
            results = response["vulnerabilities"]
            for i in results:
                fh.write(json.dumps(i["cve"]))
                fh.write("\n")

            progress_bar.update(response["resultsPerPage"])

    progress_bar.close()

#    json_out = json.dumps(results, indent=2)
#    with open("nvd-out.json", "w") as fh:
#      fh.write(json_out)

if __name__ == "__main__":
    main()