#!/usr/bin/env python3

import os

import json

from datetime import datetime, timedelta

import logging

from urllib import request, parse

import time

logging.basicConfig(

level=logging.INFO,

format="%(asctime)s: %(levelname)s: %(message)s")

BASE_URL = os.getenv('SCANNER_NVD_URL', "https://services.nvd.nist.gov/rest/json/cves/2.0")

API_KEY = os.getenv('SCANNER_NVD_API_KEY', None)

BACKOFF_TABLE = [15, 30, 60, 120, 240, 300]

class ApiLoader:

def __init__(self):

pass

@staticmethod

def get_dates(year, interval):

# interval should be a factor of 12

if 12 % interval != 0:

raise ValueError("Interval must be a factor of 12.")

current_date = datetime.now()

for month in range(1, 13, interval):

start = datetime(year, month, 1)

# return if start date is in the future

if start > current_date:

return

# Handling for months greater than 12 (for intervals like 4, 6 months, etc.)

end_month = month + interval

end_year = year

if end_month > 12:

end_month -= 12

end_year += 1

end = datetime(end_year, end_month, 1) - timedelta(days=1)

yield start.strftime("%Y-%m-%dT00:00:00.000"), end.strftime("%Y-%m-%dT23:59:59.999")

@staticmethod

def fetch_data(start, end):

index, total = 0, 0

logging.info(f"Fetching page from {start} to {end}")

max_retries = 10

attempt = 0

while total == 0 or index < total:

try:

params = {

'noRejected': '',

'startIndex': index,

'pubStartDate': start,

'pubEndDate': end

}

encoded_params = parse.urlencode(params)

url = f"{BASE_URL}?{encoded_params}"

req = request.Request(url)

req.headers = {'apiKey': API_KEY}

with request.urlopen(req, timeout=120) as response:

data = json.loads(response.read().decode())

total = data['totalResults']

logging.info(f"Fetched page at index {index} (out of {total} total items)")

for item in data['vulnerabilities']:

yield {"cve": item['cve']}

index += data['resultsPerPage']

attempt = 0

except Exception as e:

if attempt < max_retries:

backoff = BACKOFF_TABLE[min(attempt, len(BACKOFF_TABLE) - 1)]

logging.warning(f"Failed to download page at index {index} (attempt {attempt + 1}/{max_retries + 1}): {e}")

logging.info(f"Retrying in {backoff} seconds...")

time.sleep(backoff)

attempt += 1

else:

raise

time.sleep(2)

@staticmethod

def start(dirpath):

try:

for year in range(2002, datetime.now().year + 1):

logging.info(f"Fetching year {year}")

vulnerabilities = []

for start, end in ApiLoader.get_dates(year, 3): # fetch quarterly

vulnerabilities.extend(ApiLoader.fetch_data(start, end))

file_path = os.path.join(dirpath, f"{year}.nvd.json")

with open(file_path, "w") as file:

for start, end in ApiLoader.get_dates(year, 3): # fetch quarterly

for vulnerability in ApiLoader.fetch_data(start, end):

file.write(json.dumps(vulnerability) + "\n")

file_size = os.path.getsize(file_path) # Get the size of the file

logging.info(f"Size of {file_path}: {file_size} bytes")

except Exception as e:

logging.error(f"An error occurred: {e}")

raise # Re-raise the exception

def main():

import argparse

parser = argparse.ArgumentParser()

if API_KEY is None:

raise ValueError("API_KEY is not set. Please provide a valid API key.")

def dirpath(s):

if os.path.isdir(s):

return s

raise ValueError

parser.add_argument(

'dirpath',

help="Path to directory where the NVD data will be saved.",

type=dirpath)

args = parser.parse_args()

ApiLoader.start(args.dirpath)

if __name__ == "__main__":

main()