# from splunk.persistconn.application import PersistentServerConnectionApplication

import os
import sys
import time
import datetime
import gzip
import urllib3
import json


def download_nvd(year=2020, proxysettings=None, data_path="../lookups"):
    """Download and extract the json data from the NVD website."""
    url = "https://nvd.nist.gov/feeds/json/cve/1.1/nvdcve-1.1-{}.json.gz".format(year)
    
    if proxysettings is None:
        http = urllib3.PoolManager()
        response = http.request('GET', url, preload_content=False)

    elif proxysetting is not None:
        # get proxy setting configuration
        user = proxysettings['user']
        password = proxysettings['password']
        proxy_ip = proxysettings['ip']
        proxy_port = proxysettings['port']

        default_headers = urllib3.make_headers(
                             proxy_basic_auth='{}:{}'.format(user, password))
        proxy_addr = 'http://{}:{}'.format(proxy_ip, proxy_port)
        proxy = urllib3.ProxyManager(proxy_url=proxy_addr, proxy_headers=default_headers, 
                                     cert_reqs='CERT_NONE', assert_hostname=False)
        response = proxy.request('GET', url, preload_content=False)

    # Create the folder if it is missing
    if not os.path.isdir(data_path):
        os.mkdir(data_path)

    # Save zipped file to drive
    zip_file_path = os.path.join(data_path, "nvdcve-1.1-{}.json.gz".format(year))
    with open(zip_file_path, "wb") as zipfile:
        zipfile.write(response.data)

    # Extract the zipped file to text/json
    json_file_path = os.path.join(data_path, "nvdcve-1.1-{}.json".format(year))
    with gzip.GzipFile(zip_file_path, 'rb') as decompressed_file:
        json_string = decompressed_file.read()
        
    # Save the extracted file
    with open(json_file_path, 'wb') as json_file:
        json_file.write(json_string)

    
def convert_json_to_list(data_path, year):
    """Parse the vulnerability json to a python list."""
    json_file_path = os.path.join(data_path, "nvdcve-1.1-{}.json".format(year))

    with open(json_file_path, "r") as cve_json:
        cve_dict = json.load(cve_json)

    vul_list = []
    for cve_index in range(len(cve_dict['CVE_Items'])):
        # Extract the CVE number of the vulnerability

        cve_id = cve_dict['CVE_Items'][cve_index]['cve']['CVE_data_meta']['ID']
        description = cve_dict['CVE_Items'][cve_index]['cve']['description']['description_data'][0]['value']
        published_date = cve_dict['CVE_Items'][cve_index]['publishedDate']
        # Extract the score if available
        try:
            cvss = cve_dict['CVE_Items'][cve_index]['impact']['baseMetricV3']['cvssV3']['baseScore']
        except:
            cvss = "N/A"

       
        # Not every CVE id has a machine-readable config (cpe_match) entry.
        # Supress exceptions with try.. except      
        try:
            entries = cve_dict['CVE_Items'][cve_index]['configurations']['nodes']
            versions = set()
            for config_index in range(len(entries)):
                
                # Extract list of all vulnerable configurations (for 'OR' connector)
                vulnerable_configs = []
                if entries[config_index]['operator'] == 'OR':
                    for cpe in entries[0]['cpe_match']:
                        if cpe['vulnerable'] is True:
                            try:
                                max_vers = cpe['versionEndExcluding']
                            except:
                                max_vers = None
                            # aCPE string
                            cpe_string = cpe['cpe23Uri']
                            vulnerable_configs += [(cpe_string,  max_vers)]               
                            
                else:
                    for config in entries:
                        for cpe in config['children'][0]['cpe_match']:
                            if cpe['vulnerable'] is True:
                                try:
                                    max_vers = cpe['versionEndExcluding']
                                except:
                                    max_vers = None
                                # aCPE string
                                cpe_string = cpe['cpe23Uri']
                                vulnerable_configs += [(cpe_string,  max_vers)]
                
                # Remove duplicates
                unique_vulnerable_configs = list(dict.fromkeys(vulnerable_configs))
                
                for cpe, max_vers in unique_vulnerable_configs:
                    vendor = cpe.split(':')[3]
                    name = cpe.split(':')[4]
                    version = cpe.split(':')[5]
                    if version == '*':
                        if max_vers is not None:
                            version = max_vers
                    versions.add(version)
                    
            versions_string = '|'.join(versions)

        # Put in NaN in for name, version, vendor if unavailable in machine-readable cpe entries
        except:
            vendor = "N/A"
            name = "N/A"
            versions_string = "N/A"
            published_date = "N/A"

        vul_list += [[cve_id, vendor, name, versions_string, description, cvss, published_date]]

    return vul_list


def export_list_to_csv(vul_list, data_path, year):
    """Save the vulnerability list to a csv file."""

    csv_file_path = os.path.join(data_path, "nvdcve-1.1-{}.csv".format(year))
    # Format CVE table as comma-seperated multiline string
    csv_string = "cve_id, vendor, name, version, description, cvss, published_date\n"            

    if sys.version_info.major == 3:
        print('Python 3 detected.')
        for line in vul_list:
            for word in line:
                sanitized_word = str(word).replace(",", " ").replace('"','')
                csv_string += '"' + sanitized_word + '", '
        csv_string += "\n"
        print(csv_string[0:500])
        print('Exporting to csv.')
        # Save CVE table to a comma-seperated values (csv) file
        with open(csv_file_path, "w") as csv_file:
            csv_file.write(csv_string)

    elif sys.version_info.major == 2:
        print('Python 2 detected.')
        for line in vul_list:
            for word in line:
                if isinstance(word, unicode):
                    ascii_word = word.encode('ascii', 'ignore')
                elif isinstance(word, float):
                    ascii_word = str(word)
                else:
                    ascii_word = "N/A"

                sanitized_word = ascii_word.replace(",", " ").replace('"','')
                csv_string += '"' + sanitized_word + '", '
        csv_string += "\n"

        # Save CVE table to a comma-seperated values (csv) file
        csv_file_path = json_file_path.replace("json", "csv")
        print('Exporting to csv.')
        with open(csv_file_path, "w") as csv_file:
            csv_file.write(csv_string.encode('ascii'))

    else:
        raise ValueError('Invalid python version detected.')

if __name__ == "__main__":
    if False:
        proxysettings['user'] = 'proxyusername'
        proxysettings['password'] = 'proxypassword'
        proxysettings['ip'] = 'proxyip'
        proxysettings['port'] = 'port'
    else:
        proxysettings = False

    year = "2020"
    relative_data_path = "data"
    script_path, scriptname  = os.path.split(os.path.realpath(__file__))
    data_path = os.path.join(script_path, relative_data_path)

    download_nvd(data_path=data_path, year=year, proxysettings=proxysettings)
    vul_list = convert_json_to_list(data_path=data_path, year=year)
    export_list_to_csv(vul_list, data_path=data_path, year=year)
    print("Done.")
