# Important note:
# This product uses the NVD API but is not endorsed or certified by the NVD.

require cve-update-native.inc

NVDCVE_URL ?= "https://services.nvd.nist.gov/rest/json/cves/2.0"

# If you have a NVD API key (https://nvd.nist.gov/developers/request-an-api-key)
# then setting this to get higher rate limits.
NVDCVE_API_KEY ?= ""

# CVE database incremental update age threshold, in seconds. If the database is
# older than this threshold, do a full re-download, else, do an incremental
# update. By default: the maximum allowed value from NVD: 120 days (120*24*60*60)
# Use 0 to force a full download.
CVE_DB_INCR_UPDATE_AGE_THRES ?= "10368000"

# Number of attempts for each http query to nvd server before giving up
CVE_DB_UPDATE_ATTEMPTS ?= "5"

def nvd_request_wait(attempt, min_wait):
    return min(((2 * attempt) + min_wait), 30)

def nvd_request_next(url, attempts, api_key, args, min_wait):
    """
    Request next part of the NVD database
    NVD API documentation: https://nvd.nist.gov/developers/vulnerabilities
    """

    import urllib.request
    import urllib.parse
    import gzip
    import http
    import time

    request = urllib.request.Request(url + "?" + urllib.parse.urlencode(args))
    if api_key:
        request.add_header("apiKey", api_key)
    bb.note("Requesting %s" % request.full_url)

    for attempt in range(attempts):
        try:
            r = urllib.request.urlopen(request)

            if (r.headers['content-encoding'] == 'gzip'):
                buf = r.read()
                raw_data = gzip.decompress(buf)
            else:
                raw_data = r.read().decode("utf-8")

            r.close()

        except Exception as e:
            wait_time = nvd_request_wait(attempt, min_wait)
            bb.note("CVE database: received error (%s)" % (e))
            bb.note("CVE database: retrying download after %d seconds. attempted (%d/%d)" % (wait_time, attempt+1, attempts))
            time.sleep(wait_time)
            pass
        else:
            return raw_data
    else:
        # We failed at all attempts
        return None

def update_db_file(db_tmp_file, d, database_time):
    """
    Update the given database file
    """
    import bb.progress
    import bb.utils
    import datetime
    import sqlite3
    import json

    pre_update_utc_timestamp = datetime.datetime.now().astimezone(tz=datetime.timezone.utc)

    # Connect to database
    conn = sqlite3.connect(db_tmp_file)
    initialize_db(conn)

    req_args = {'startIndex': 0}

    incr_update_threshold = int(d.getVar("CVE_DB_INCR_UPDATE_AGE_THRES"))
    if database_time != 0:
        database_date = datetime.datetime.fromtimestamp(database_time, tz=datetime.timezone.utc)
        today_date = datetime.datetime.now(tz=datetime.timezone.utc)
        delta = today_date - database_date
        if incr_update_threshold == 0:
            bb.note("CVE database: forced full update")
        elif delta < datetime.timedelta(seconds=incr_update_threshold):
            bb.note("CVE database: performing partial update")
            # The maximum range for time is 120 days
            if delta > datetime.timedelta(days=120):
                bb.error("CVE database: Trying to do an incremental update on a larger than supported range")
            req_args['lastModStartDate'] = database_date.isoformat()
            req_args['lastModEndDate'] = today_date.isoformat()
        else:
            bb.note("CVE database: file too old, forcing a full update")
    else:
        bb.note("CVE database: no preexisting database, do a full download")

    with bb.progress.ProgressHandler(d) as ph, open(os.path.join(d.getVar("TMPDIR"), 'cve_check'), 'a') as cve_f:

        bb.note("Updating entries")
        index = 0
        url = d.getVar("NVDCVE_URL")
        api_key = d.getVar("NVDCVE_API_KEY") or None
        attempts = int(d.getVar("CVE_DB_UPDATE_ATTEMPTS"))

        # Recommended by NVD
        wait_time = 6
        if api_key:
            wait_time = 2

        while True:
            req_args['startIndex'] = index
            raw_data = nvd_request_next(url, attempts, api_key, req_args, wait_time)
            if raw_data is None:
                # We haven't managed to download data
                return False

            # hack for json5 style responses
            if raw_data[-3:] == ',]}':
                bb.note("Removing trailing ',' from nvd response")
                raw_data = raw_data[:-3] + ']}'

            data = json.loads(raw_data)

            index = data["startIndex"]
            total = data["totalResults"]
            per_page = data["resultsPerPage"]
            bb.note("Got %d entries" % per_page)
            for cve in data["vulnerabilities"]:
                update_db(conn, cve)

            index += per_page
            ph.update((float(index) / (total+1)) * 100)
            if index >= total:
                break

            # Recommended by NVD
            time.sleep(wait_time)

        # Update success, set the date to cve_check file.
        cve_f.write('CVE database update : %s\n\n' % datetime.date.today())
        conn.execute("insert into MTIME values (?)", [pre_update_utc_timestamp.isoformat()]).close()

    conn.commit()
    conn.close()
    return True


def parse_node_and_insert(conn, node, cveId):

    def cpe_generator():
        for cpe in node.get('cpeMatch', ()):
            if not cpe['vulnerable']:
                return
            cpe23 = cpe.get('criteria')
            if not cpe23:
                return
            cpe23 = cpe23.split(':')
            if len(cpe23) < 6:
                return
            vendor = cpe23[3]
            product = cpe23[4]
            version = cpe23[5]

            if cpe23[6] == '*' or cpe23[6] == '-':
                version_suffix = ""
            else:
                version_suffix = "_" + cpe23[6]

            if version != '*' and version != '-':
                # Version is defined, this is a '=' match
                yield [cveId, vendor, product, version + version_suffix, '=', '', '']
            elif version == '-':
                # no version information is available
                yield [cveId, vendor, product, version, '', '', '']
            else:
                # Parse start version, end version and operators
                op_start = ''
                op_end = ''
                v_start = ''
                v_end = ''

                if 'versionStartIncluding' in cpe:
                    op_start = '>='
                    v_start = cpe['versionStartIncluding']

                if 'versionStartExcluding' in cpe:
                    op_start = '>'
                    v_start = cpe['versionStartExcluding']

                if 'versionEndIncluding' in cpe:
                    op_end = '<='
                    v_end = cpe['versionEndIncluding']

                if 'versionEndExcluding' in cpe:
                    op_end = '<'
                    v_end = cpe['versionEndExcluding']

                if op_start or op_end or v_start or v_end:
                    yield [cveId, vendor, product, v_start, op_start, v_end, op_end]
                else:
                    # This is no version information, expressed differently.
                    # Save processing by representing as -.
                    yield [cveId, vendor, product, '-', '', '', '']

    conn.executemany("insert into PRODUCTS values (?, ?, ?, ?, ?, ?, ?)", cpe_generator()).close()

def update_db(conn, elt):
    """
    Update a single entry in the on-disk database
    """

    accessVector = None
    vectorString = None
    cveId = elt['cve']['id']
    if elt['cve'].get('vulnStatus') == "Rejected":
        c = conn.cursor()
        c.execute("delete from PRODUCTS where ID = ?;", [cveId])
        c.execute("delete from NVD where ID = ?;", [cveId])
        c.close()
        return
    cveDesc = ""
    for desc in elt['cve']['descriptions']:
        if desc['lang'] == 'en':
            cveDesc = desc['value']
    date = elt['cve']['lastModified']
    try:
        accessVector = elt['cve']['metrics']['cvssMetricV2'][0]['cvssData']['accessVector']
        vectorString = elt['cve']['metrics']['cvssMetricV2'][0]['cvssData']['vectorString']
        cvssv2 = elt['cve']['metrics']['cvssMetricV2'][0]['cvssData']['baseScore']
    except KeyError:
        cvssv2 = 0.0
    cvssv3 = None
    try:
        accessVector = accessVector or elt['cve']['metrics']['cvssMetricV30'][0]['cvssData']['attackVector']
        vectorString = vectorString or elt['cve']['metrics']['cvssMetricV30'][0]['cvssData']['vectorString']
        cvssv3 = elt['cve']['metrics']['cvssMetricV30'][0]['cvssData']['baseScore']
    except KeyError:
        pass
    try:
        accessVector = accessVector or elt['cve']['metrics']['cvssMetricV31'][0]['cvssData']['attackVector']
        vectorString = vectorString or elt['cve']['metrics']['cvssMetricV31'][0]['cvssData']['vectorString']
        cvssv3 = cvssv3 or elt['cve']['metrics']['cvssMetricV31'][0]['cvssData']['baseScore']
    except KeyError:
        pass
    cvssv3 = cvssv3 or 0.0
    try:
        accessVector = accessVector or elt['cve']['metrics']['cvssMetricV40'][0]['cvssData']['attackVector']
        vectorString = vectorString or elt['cve']['metrics']['cvssMetricV40'][0]['cvssData']['vectorString']
        cvssv4 = elt['cve']['metrics']['cvssMetricV40'][0]['cvssData']['baseScore']
    except KeyError:
        cvssv4 = 0.0
    accessVector = accessVector or "UNKNOWN"
    vectorString = vectorString or "UNKNOWN"

    conn.execute("insert or replace into NVD values (?, ?, ?, ?, ?, ?, ?, ?)",
                [cveId, cveDesc, cvssv2, cvssv3, cvssv4, date, accessVector, vectorString]).close()

    try:
        # Remove any pre-existing CVE configuration. Even for partial database
        # update, those will be repopulated. This ensures that old
        # configuration is not kept for an updated CVE.
        conn.execute("delete from PRODUCTS where ID = ?", [cveId]).close()
        for config in elt['cve']['configurations']:
            # This is suboptimal as it doesn't handle AND/OR and negate, but is better than nothing
            for node in config["nodes"]:
                parse_node_and_insert(conn, node, cveId)
    except KeyError:
        bb.note("CVE %s has no configurations" % cveId)
