#!/usr/bin/env python3

import sys
import math
import json
import time
from pymongo import MongoClient
import requests
from pprint import pformat
from pprint import pprint


SINGLE_ENDPOINT     = "https://services.nvd.nist.gov/rest/json/cve/1.0"
COLLECTION_ENDPOINT = "https://services.nvd.nist.gov/rest/json/cves/1.0"

client = MongoClient('mongodb://localhost:27017/')
db     = client.cves # The database is called cves.
cves   = db.cves     # The database has only one collection with cves.
    
def getSingleCVE(cve):
    url = "/".join([SINGLE_ENDPOINT, cve])
    res = requests.get(url)
    
    return res.json()

def getCVECollection(params):
    res = requests.get(COLLECTION_ENDPOINT, params=params)
    print(f"Request URL is {res.url}")
    print(f"Status code: {res.status_code}")

    return res.json()

def main():
    pageSize = 1024
    params = {
        "startIndex": 0,
        "resultsPerPage": 0,
        "pubStartDate": "2010-01-01T00:00:00:000 UTC-00:00",
        "pubEndDate":   "2021-01-01T00:00:00:000 UTC-00:00",
    }

    # Get total number of results.
    cveRawCollection = getCVECollection(params)
    totalResults = cveRawCollection["totalResults"]
    print(f"Total number of results is {totalResults}")

    pageCount = math.ceil(totalResults / pageSize)
    print(f"Number of pages is equal to {pageCount}")

    time.sleep(10)

    print(f"Gatering info from NVD...")

    start = time.time()
    params["resultsPerPage"] = pageSize
    for pageNo in range(pageCount):
        params["startIndex"] = pageNo * pageSize

        parsedCVEList = []
        cveRawCollection = getCVECollection(params)
        for cve in cveRawCollection["result"]["CVE_Items"]:
            id = cve["cve"]["CVE_data_meta"]["ID"].lower()
            parsedCVEList.append({
                "_id": id,
                "nvd": cve,
            })
        # cves.insert_many(parsedCVEList)
    
    end = time.time()
    elapsed = end - start
    print(f"Done. Time elpased: {elapsed}")

    # with open("nvd_cves_collection_2010_2020.json", "w") as f:
    #     json.dump(parsedCVEList, f)


if __name__ == "__main__":
    main()
