UPDATED: Multithread-based

Author: Ruskonert

nvdnist.py

@@ -9,7 +9,8 @@
 import os
 import signal
 import time
-import copy
+import threading
+
 from multiprocessing import Process
 from database import TableColumn
 from threading import Thread
@@ -29,6 +30,7 @@ def __init__(self, process):
         self._process = None
         self.kill = False
         self.process = process
+        self._lock = threading.Lock()
 
     def _pobject(self):
         self._entry = len(self._wpList)
@@ -61,20 +63,21 @@ def _pobject(self):
             if len(self._wpList) == 0:
                 self.kill = True
 
-    def _processFunction(self):
+    def _processFunction(self, delay=6.0):
         while True:
             sys.stdout.write("\rCurrent progress: %d / %d" % (self._entry - len(self._wpList), self._entry))
             sys.stdout.flush()
             if self.kill:
+                self._lock.release()
                 break
-            time.sleep(5)
+            time.sleep(delay)
 
     def start(self):
         self._coreThread = Thread(target=self._pobject)
         self._coreThread.start()
-
         self._process = Thread(target=self._processFunction)
         self._process.start()
+        self._lock.acquire()
 
     def add_thread(self, thread):
         self._wpList.append(thread)
@@ -97,39 +100,21 @@ def workstation(sql, tablename, jsonElement):
             cve_data_version = str(jsonElement['CVE_data_version'])
             reference = ""
             description = ""
-            cpe = ""
             problemtype = ""
 
             for r in base.find_key(jsonElement, r"reference_.+$"):
-                reference += jsonElement[r] + " || "
-            reference = reference[:-4]
+                reference += jsonElement[r] + " | "
+            reference = reference[:-3]
 
             for d in base.find_key(jsonElement, r"desc_(.+)"):
-                description += jsonElement[d] + " || "
-            description = description[:-4]
+                description += jsonElement[d] + " | "
+            description = description[:-3]
 
             for d in base.find_key(jsonElement, r"cwe_\d"):
-                problemtype += jsonElement[d] + " || "
-            problemtype = problemtype[:-4]
-
-            cpe_list = []
-
-            for c in base.find_key(jsonElement, r"cpe_\d_vulnderable"):
-                cpe_list.append(str(jsonElement[c]) + "|")
-
-                
-            for index, c in enumerate(base.find_key(jsonElement, r"cpe_\d_cpe22uri")):
-                cpe_list[index] += (str(jsonElement[c])) + "|"
-
-
-            for index, c in enumerate(base.find_key(jsonElement, r"cpe_\d_cpe23uri")):
-                cpe_list[index] += (str(jsonElement[c]))
-
-            for c in cpe_list:
-                cpe += c + " || "
-            cpe = cpe[:-4]
+                problemtype += jsonElement[d] + " | "
+            problemtype = problemtype[:-3]
 
-                # for version 2
+            # for version 2
             baseMetricV2_version = str(jsonElement.get('baseMetricV2.version', ''))
             baseMetricV2_vectorString = str(jsonElement.get('baseMetricV2.vectorString', ''))
             baseMetricV2_accessVector = str(jsonElement.get('baseMetricV2.accessVector', ''))
@@ -167,43 +152,76 @@ def workstation(sql, tablename, jsonElement):
             publishedDate = str(jsonElement.get('publishedDate', ''))
             lastModifiedDate = str(jsonElement.get('lastModifiedDate', ''))
 
-            affect = ""
+            cpe_list = []
+            r = base.find_key(jsonElement, r"cpe_\d_vulnderable")
+            if len(r) != 0:
+                for index, c in enumerate(r):
+                    cpe_list.append(str(jsonElement[c]) + "||")
+                    
+                for index, c in enumerate(base.find_key(jsonElement, r"cpe_\d_cpe22uri")):
+                    cpe_list[index] += (str(jsonElement[c])) + "||"
+
+
+                for index, c in enumerate(base.find_key(jsonElement, r"cpe_\d_cpe23uri")):
+                    cpe_list[index] += (str(jsonElement[c]))
+
+            # each of affect data.
             for index, c in enumerate(base.find_key(jsonElement, r"affect_\d")):
-                affect += jsonElement[c] + " || "
-            affect = affect[:-4]
-
-            a = None
-            r = None
-            d = None
-            c = None
-            query_argument = None
-            query_format = None
-            query_argument = None
-            query_header = None
-            variable_name = None
-
-            v = dict(vars())
-            
-            query_header = "("
-            query_format = "("
-            query_argument = [] 
-
-            # automatic insert value
-            for variable_name in v.keys():
-                if isinstance(v[variable_name], str) and not variable_name == "tablename":
-                    query_header += variable_name + ','
-                    query_format += "%s,"
-                    query_argument.append(v[variable_name])
-            query_header = query_header[:-1]
-            query_header += ")"
-            query_format = query_format[:-1]
-            query_format += ")"
-            sucessful = base.send_query("INSERT INTO {table_name} {header} VALUES {format}".format(table_name=tablename, header=query_header, format=query_format), tuple(query_argument), False, sql.execption_list)
+                r = jsonElement[c].split('|')
+                affect_vendor_name = r[0]
+                affect_product_name = r[1]
+                affect_version_value = r[2]
+
+                nodes_cpe_vulnerable = None
+                nodes_cpe_cpe22uri = None
+                nodes_cpe_cpe23uri = None
+                
+                if len(cpe_list) != 0:
+                    for c in cpe_list:
+                        r = c.split('||')
+                        nodes_cpe_vulnerable = r[0]
+                        nodes_cpe_cpe22uri = r[1]
+                        nodes_cpe_cpe23uri = r[2]
+                        a = None
+                        r = None
+                        d = None
+                        c = None
+                        v = dict(vars())
+                        execute(base, sql, tablename, v)
+                else:
+                    a = None
+                    r = None
+                    d = None
+                    c = None
+                    v = dict(vars())
+                    execute(base, sql, tablename, v)
             base.disconnect()
         except KeyboardInterrupt:
-            sys.exit()
+            sys.exit(0)
+    
         except SystemExit:
-            sys.exit()
+            sys.exit(0)
+
+def execute(base, sql, tablename, v):
+    query_header = "("
+    query_format = "("
+    query_argument = [] 
+
+    # automatic insert value
+    for variable_name in v.keys():
+        if isinstance(v[variable_name], str) and not variable_name == "tablename":
+            query_header += variable_name + ','
+            query_format += "%s,"
+            query_argument.append(v[variable_name])
+
+    query_header = query_header[:-1]
+    query_header += ")"
+
+    query_format = query_format[:-1]
+    query_format += ")"
+
+    base.send_query("INSERT INTO {table_name} {header} VALUES {format}".format(table_name=tablename, header=query_header, format=query_format), tuple(query_argument), False, sql.execption_list)
+
 
 class Nvdnist(CVEDatabase):
     def __init__(self, mysql_base):
@@ -331,20 +349,30 @@ def load(self, obj):
         return raw_file, collection
 
     def default_table_column(self):
+        # There's no meaning about variable name.
+        # Because it reads automantually only value.
         a2 = TableColumn("year", "VARCHAR(5)")
-
         a3 = TableColumn("ID", "VARCHAR(24)")
         a4 = TableColumn('ASSIGNER', "VARCHAR(15)")
         a5 = TableColumn("data_type", 'VARCHAR(5)')
         a6 = TableColumn("data_format", "VARCHAR(10)")
         a7 = TableColumn("data_version", "VARCHAR(5)")
-        a8 = TableColumn("affect", "TEXT")
-        aa8 = TableColumn("problemtype", "VARCHAR(25)")
+
+        a8 = TableColumn("affect_vendor_name", "TEXT")
+        aa8 = TableColumn("affect_product_name", "TEXT")
+        ab8 = TableColumn("affect_version_value", "TEXT")
+        
+        ac8 = TableColumn("problemtype", "VARCHAR(25)")
         a9 = TableColumn("description", "TEXT")
         a10 = TableColumn("reference", "TEXT")
         a11 = TableColumn("cve_data_version", "VARCHAR(5)")
-        a12 = TableColumn("cpe", "TEXT")
-        
+
+        cpe1 = TableColumn("nodes_operator", "VARCHAR(5)")
+        cpe2 = TableColumn("nodes_cpe", "TEXT")
+        cpe3 = TableColumn("nodes_cpe_vulnerable", "VARCHAR(5)")
+        cpe4 = TableColumn("nodes_cpe_cpe22uri", "TEXT")
+        cpe5 = TableColumn("nodes_cpe_cpe23uri", "TEXT")
+
         a13 = TableColumn('baseMetricV2_version', "VARCHAR(10)")
         a14 = TableColumn('baseMetricV2_vectorString', "VARCHAR(50)")
         a15 = TableColumn('baseMetricV2_confidentialityImpact', "VARCHAR(10)")
@@ -392,19 +420,21 @@ def find_key(self, element, pattern):
         return key
 
     def table_insert(self, tablename, jsonCollection):
-        try:
-            print("MYSQL inserting thread now started")
-            threadlist = ThreadList(thread_count)
-            print("Inserting value ...")
-            for jsonElement in jsonCollection:
-                p = Thread(target=workstation, args=(self, tablename, jsonElement))
-                threadlist.add_thread(p)
-            threadlist.start()
-            while threadlist.kill == False:
-                time.sleep(0.1)
-        except (KeyboardInterrupt, SystemExit):
-            sys.exit(-1)
+        threadlist = ThreadList(thread_count)
+        print("Inserting value ...")
+            
+        # Singlethreading-based
+        #p = Thread(target=self._processFunction, args=())
+        #for jsonElement in jsonCollection:
+        #    workstation(self, tablename, jsonElement)
 
+        # Multithreading-based
+        print("MYSQL inserting thread now started")
+        for jsonElement in jsonCollection:
+            p = Thread(target=workstation, args=(self, tablename, jsonElement))
+            threadlist.add_thread(p)
+        threadlist.start()
+        threadlist._process.join()
         print("\nAll data has been inserted.")
 
 
@@ -606,7 +636,7 @@ def affect_vendor_data(self, cve_affects):
                     product_name = product_info['product_name']
                     for version in product_info['version']['version_data']:
                         version_value = version
-                        product_flag.append(product_name + ":" + version_value.get('version_value', 'invaild')) 
+                        product_flag.append(product_name + "|" + version_value.get('version_value', 'invaild')) 
                 affect = Affect(vendor_name, product_flag)
                 vendor_list.append(affect)
             except:
@@ -709,7 +739,7 @@ def __init__(self, vendor_name, product_flag):
 
     def serialize(self, jsonDefault, o):
         for index, flag in enumerate(self.product_flag):
-            jsonDefault['affect_{n}'.format(n=index)] = str(self.vendor_name + ":" + flag)
+            jsonDefault['affect_{n}'.format(n=index)] = str(self.vendor_name + "|" + flag)
 
 class DescriptionType:
     __dbcolumn__ = "Description"
@@ -718,14 +748,14 @@ def __init__(self, lang, value):
         self.value = value
 
     def serialize(self, jsonDefault, o):
-            jsonDefault["desc_" + str(id(self))] = self.lang + "|" + self.value
+            jsonDefault["desc_" + str(id(self))] = self.lang + ", " + self.value
 
 class ProblemType(DescriptionType):
     def __init__(self, lang, value):
         super().__init__(lang, value)
 
     def serialize(self, jsonDefault, o):
-        jsonDefault["cwe_" + str(id(self))] = self.lang + "|" + self.value
+        jsonDefault["cwe_" + str(id(self))] = self.lang + ", " + self.value
 
 class Reference:
     def __init__(self, url, name, refsource):
@@ -734,4 +764,4 @@ def __init__(self, url, name, refsource):
         self.refsource = refsource
 
     def serialize(self, jsonDefault, o):
-            jsonDefault["reference_" + self.name] = self.url + "|" + self.refsource
+            jsonDefault["reference_" + self.name] = self.url + ", " + self.refsource