cve5-parse.py

#!/usr/bin/env python3

import sys
import json
import os
import urllib.request
import elasticsearch
import elasticsearch.helpers
from elasticsearch import Elasticsearch

if 'ESURL' not in os.environ:
    es_url = "http://localhost:9200"
else:
    es_url = os.environ['ESURL']

if 'ESCERT' in os.environ:
    cert = os.environ['ESCERT']
else:
    cert = None

if cert:
    es = Elasticsearch([es_url], ca_certs=cert)
else:
    es = Elasticsearch([es_url])

class CVE:

    def __init__(self, update_rate = 1000):
        self.ids = []
        self.current = -1
        self.rh_data = None
        self.update_rate = update_rate

    def add(self, i):
        # some of these don't exist, just give up if it fails
        i['id'] = i["cveMetadata"]['cveId']
        i['year'] = int(i['id'].split('-')[1])
        i['just_id'] = int(i['id'].split('-')[2])
        print(i['id'])

        if 'x_legacyV4Record' in i["containers"]["cna"]:
            del(i["containers"]["cna"]["x_legacyV4Record"])
        if 'x_generator' in i["containers"]["cna"]:
            del(i["containers"]["cna"]["x_generator"])

        try:
            i['description_len'] = len(i["containers"]["cna"]["descriptions"][0]["value"])
        except:
            pass

        cve_bulk = {
                    "_op_type": "update",
                    "_index":   "cve5-index",
                    "_id":      i["id"],
                    "doc_as_upsert": True,
                    "doc":  i
                   }

        self.ids.append(cve_bulk)
        self.__check_update()

    def __check_update(self, force = False):

        if self.update_rate == 1:
            resp = es.index(index="cve5-index", id=self.ids[0]["_id"], document=self.ids[0]["doc"])
            self.ids = []
            self.current = -1

        elif force or len(self) > self.update_rate:

            for ok, item in elasticsearch.helpers.streaming_bulk(es, self, max_retries=2):
                if not ok:
                    print("ERROR:")
                    print(item)

            self.ids = []
            self.current = -1

    def done(self):
        self.__check_update(True)

    def __next__(self):
        "Handle a call to next()"

        self.current = self.current + 1
        if self.current >= len(self.ids):
            raise StopIteration

        return self.ids[self.current]

    def __iter__(self):
        return self

    def __len__(self):
        return len(self.ids)

def main():

    if len(sys.argv) > 1:
        input_dir = sys.argv[1]
    else:
        print("Usage: %s <cve5 dir>" % (sys.argv[0]))
        sys.exit(1)

    # First let's see if the index exists
    if es.indices.exists(index='cve5-index') is False:
        # We have to create it and add a mapping
        # Mapping is busted
        #fh = open('cve-index-json-mapping.json')
        #mapping = json.load(fh)
        #es.indices.create(index='cve-index', mappings=mapping["mappings"], settings=mapping["settings"])
        es.indices.create(index='cve5-index')

    #the_cves = CVE(update_rate = 1)
    the_cves = CVE()
    for root, dirs, files in os.walk(input_dir):
        for name in files:
            if name.endswith(".json") and name.startswith("CVE"):
                full_path = os.path.join(root, name)
                print(full_path)
                fh = open(full_path)

                line = " ".join(fh.readlines())
                if not line: break

                json_data = json.loads(line)

                the_cves.add(json_data)

    the_cves.done()

if __name__ == "__main__":
    main()