import urllib.request
import json
import datetime
import gzip
import math
import time
import random
import xml.etree.ElementTree as ET

lastDate = datetime.datetime.now()
cpeDict = {}

def buildCpeDict():

    bS = "{http://cpe.mitre.org/dictionary/2.0}"
    aS = "{http://scap.nist.gov/schema/cpe-extension/2.3}"

    tree = ET.parse('cpe-dict.xml')
    dictRoot = tree.getroot()
    cpeItems = dictRoot.findall(bS + 'cpe-item')

    print("[*] Counting total CPEs")
    try:
        metaPull =  urllib.request.urlopen(
                        "https://services.nvd.nist.gov/rest/json/cpes/1.0")
    except Exception as e:
        print("[!] Failure: " + str(e))
        print("[!] Stopping. Check Network and system statuses.")
        exit(17)
    try:
        meta = json.loads(gzip.decompress(metaPull.read()))
    except:
        meta = json.loads(metaPull.read())
    requestsCount = int(math.ceil(meta['totalResults']/250))
    print("[*] Total Requests will be " + str(requestsCount) + " at 250 per \
page")
    time.sleep(3.5)
    for i in range(0, requestsCount-1):
        success = False
    cpeData = {}
    print("[!] Gathering data " + str(i+1) + " of 2861.")
    count = 0
    while not success:
        try:
            if isinstance(lastDate, datetime.date):
                date = lastDate.strftime('yyyy-MM-dd')
                result = urllib.request.urlopen(
                    "https://services.nvd.nist.gov/rest/json/cpes/1.0?"
                    +"startIndex="+str(i*250)+"&resultsPerPage=250&"
                    +"addOns=cves&?"
                    +"modStartDate="+date+"T00:00:00:000 UTC-05:00")
            else:
                result = urllib.request.urlopen(
                    "https://services.nvd.nist.gov/rest/json/cpes/1.0?"
                    +"startIndex="+str(i*250)+"&resultsPerPage=250&"
                    +"addOns=cves")
            parse = result.read().decode('utf-8')
            cpeJson = json.loads(parse)
            for cpe in cpeJson['result']['cpes']:
                cpeString = cpe['cpe23Uri']
                try:
                    for item in cpeItems:
                        if item.find(aS + 'cpe23-item').attrib['name'] == cpeString:
                            cpeKey = item.attrib['name']
                            break
                except: 
                    cpeKey = str(cpe['titles'][0]['title'])

                cpeData.update({cpeKey: {"cpe23Uri": cpe['cpe23Uri'], "vulns":  
                                        cpe['vulnerabilities']}})
            with open("tmp/cpeDict"+str(i)+".json", "w") as cpeD:
                json.dump(cpeData, cpeD, indent=4, sort_keys=True)
            success = True
        except:
            if count%3 != 0:
                print("\t[!] Trying again after 5 minutes")
                time.sleep(300)
            else: 
                print("\t[!] Probably in the Box. Waiting 15 minutes.")
                time.sleep(500)
            success = False
            
    print("[!] Moving On to next file!")
    time.sleep(random.randint(2,5))

    print("[*] Writing to ./data/cpe-dict.json")
    with open("./data/cpe-dict.json", 'w') as cpeFile:
        json.dump(cpeDict, cpeFile, sort_keys=True, indent=4)
    
    print("[*] Clearing out Temporary Files.")


if __name__ == '__main__':
    try: 
        with open("cpe.meta",'w') as cpeMeta:
            json.load(cpeMeta)
            lastDate = cpeMeta['lastDate']
            json.dump(newMeta, cpeMeta, sort_keys=True, indent=4)
        buildCpeDict()
    except:
        buildCpeDict()
        with open("cpe.meta",'w') as cpeMeta:
            newMeta = {}
            newMeta['lastDate'] = str(datetime.datetime.now())
            lastDate = newMeta['lastDate']
            json.dump(newMeta, cpeMeta, sort_keys=True, indent=4)