UPDATED: fixed #6 & Separate process for resolves #7

Author: Ruskonert

README.md

@@ -54,7 +54,7 @@ python cvemysql.py [-u user] [-p password] [-y year] [output_name] [--cofigurati
 `[-t tablename]` 테이블 이름을 미리 설정합니다. 기본값은 `NVD_NIST_TABLE`입니다.<br>
 `[-c content]` 리스트 컨텐츠을 선택합니다. 현재는 해당 인자를 사용하지 않아도 자동으로 `NVD_NIST`로 기본 설정되어 있습니다.<br> Available type - nvd, mitre<br>
 `[-select-path]` 외부에서 RAW 파일을 다운로드 하지 않고 이미 존재하는 해당 파일 경로로 지정합니다. `$JAVA_HOME`처럼 미리 지정된 환경 변수와 혼합하여 사용할 수 있습니다.<br>
-`[-tc thread_count]`  멀티쓰레드 갯수를 설정합니다. 기본값은 `4`입니다. 해당 값을 건드리는 것은 추천하지 않습니다.
+`[-tc thread_count]`  멀티쓰레드 갯수를 설정합니다. 입력하지 않을 경우, 자동으로 지정됩니다.<br>
 `[output_file]` 출력 타입 지정시, 추출될 파일명을 지정합니다. 출력 파일의 이름을 지정하는 것을 강력히 권장합니다. 확장자는 기재하지 않아도 되며, 만약 압축 확장자(.zip)가 붙어있다면, 추출된 파일들을 전부 압축하여 `zip` 파일로 만들어줍니다.<br>
 
 
@@ -63,7 +63,7 @@ python cvemysql.py [-u user] [-p password] [-y year] [output_name] [--cofigurati
 `[--db]` 데이터베이스 파일로 추출합니다. MySQL 기반에서는 SQL 파일이 기본 확장자입니다.<br>
 `[--csv]` 데이터베이스에 있는 데이터 값들을 CSV 파일로 추출합니다.<br>
 `[--setup-install]` 처음 실행할 때 모듈이 정상적으로 작동할 수 있도록 합니다. DEPRECATED.<br>
-
+`[--max-thread]` 사용 가능한 최대 쓰레드 수로 지정합니다. 안정성을 위해 해당 인자는 권장하지 않습니다.<br>
 ## Example
 아래는 해당 플랫폼을 사용하기 위해 입력하는 명령어의 예시들을 나열한 것입니다. 프로그램을 실행시킬 때 참조하시기 바랍니다.<br>
 예를 들어 특정 년도의 NVD-NIST CVE 리스트를 모두 가져와서 데이터베이스만 만들려면 다음과 같이 하십시오:

__unused__.py

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+import threading
+from threading import Thread
+
+class ThreadList:
+    def __init__(self, process):
+        self._entry = 0
+        self.total = 0
+        self._wpList = []
+        self.pList = []
+        self._coreThread = None
+        self._process = None
+        self.kill = False
+        self.process = process
+        self._lock = threading.Lock()
+
+    def _pobject(self):
+        self._entry = len(self._wpList)
+        self._wpList.reverse()
+
+        while self.kill == False:
+            for ps in range(0, self.process):
+                try:
+                    v = self._wpList.pop()
+                    self.pList.append(v)
+                except:
+                    if len(self._wpList) != 0:
+                        self.process = len(self._wpList) % self.process
+                        self._pobject()
+                    else:
+                        self.kill = True
+                        break
+            
+            for p in self.pList:
+                p.start()
+
+            for p in self.pList:
+                p.join()
+
+            for p in self.pList:
+                del p
+        
+            self.pList.clear()
+
+            if len(self._wpList) == 0:
+                self.kill = True
+
+    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(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)

cvemitre.py

@@ -51,6 +51,12 @@ def load(self, obj):
                     is_exist = True
                 else:
                     os.remove(raw_file)
+        
+        if obj['specific_path'] is not None:
+           raw_file =  obj['specific_path'][0]
+           print("USING: %s" % raw_file)
+           is_exist = True
+        
         if not is_exist:
             target_url = url + raw_file
             print("Downloading the resource file ...")
@@ -60,7 +66,7 @@ def load(self, obj):
         if raw_type == "xml":
             readTo = ET.parse(raw_file).getroot()
         else:
-            with open(raw_file, 'r', newline='', encoding='ISO-8859-1') as fs:
+            with open(raw_file, 'r', newline='', encoding='utf-8') as fs:
                     readTo = csv.reader(fs)
                     for line in readTo:
                         try:
@@ -124,4 +130,4 @@ def insert(self, obj):
         raise NotImplementedError("No need to execute")
 
     def export(self, export_path_name="unspecified", export_type='sql', tablename="Unknown_table", option=None, charset='utf-8'):
-        return super().export(export_path_name, export_type, tablename, option, "ISO-8859-1")
+        return super().export(export_path_name, export_type, tablename, option, charset)

cvemysql.py

@@ -33,12 +33,12 @@ def initialize(user, password, host, port, char="utf8"):
     database.connect()
     return database
 
-def configure_task(base, table_name, year, output_name=None, specific_path=None):
+def configure_task(base, table_name, year, output_name=None, specific_path=None, thread_count=None):
     """ 작업을 구성합니다. """
 
     # 작업을 위한 여러가지 환경 변수를 지정해줍니다.
     directory = {"year":year,
-     "table_name":table_name, "output_name":output_name, "specific_path": specific_path}
+     "table_name":table_name, "output_name":output_name, "specific_path": specific_path, "thread_count": thread_count}
 
     # 데이터를 분석하고 이에 따른 결과값을 가져옵니다.
     filepath, collection = base.load(directory)
@@ -77,8 +77,6 @@ def help(error=None):
     if error is not None:
         raise CommandException(error)
 
-thread_count = 8
-
 if __name__ == "__main__":
 
     # 특정 파일이 지정되지 않았을 때, 외부에서 CVE RAW 데이터에 해당하는 연도를 담을 리스트입니다.
@@ -88,21 +86,22 @@ def help(error=None):
     output_filetype = []
 
     # 프로그램 인자값, 출력 파일명
-    argv, filename =  util.disting_args(sys.argv[1:], '--year-all', '--csv', '--db')
+    argv, filename =  util.disting_args(sys.argv[1:], '--year-all', '--csv', '--db', '--max-thread')
 
     # 프로그램 환경 식별
     # 쓰레드 갯수 식별
-    thread_count = int(util.define_argument_value(argv, "-thread-count", 4, flag="-tc"))
+    tcount = util.define_argument_value(argv, "-thread-count", flag="-tc")
 
     # RAW 파일 경로 여부
     file_path = util.define_argument_value(argv, "-file", None)
 
 
     # 구성 인자값 식별
-    all_years = util.define_argument_value(argv, "--all-year", False)
+    all_years = util.define_argument_value(argv, "--year-all", False)
     all_source = util.define_argument_value(argv, "--all",  False)
     csv_output = util.define_argument_value(argv, "--csv", False)
     sql_output = util.define_argument_value(argv, "--db", False)
+    max_threaded = util.define_argument_value(argv, "--max-thread", False)
 
     # 계정 인자값 식별
     user = util.define_argument_value(argv, "-user", flag="-u")
@@ -115,6 +114,22 @@ def help(error=None):
     database_name = util.define_argument_value(argv, "-dbname", flag="-d")
     table_name = util.define_argument_value(argv, "-tablename", flag="-t")
 
+    if tcount is None:
+        if max_threaded:
+            count = os.cpu_count()
+            print("경고: 사용 가능한 모든 쓰레드를 사용합니다.")
+        else:
+            count = 4
+
+        if count < 1:
+            count = 1
+        tcount = count
+    else:
+        tcount = int(tcount)
+        if max_threaded:
+            print("경고: 쓰레드 수가 이미 결정되었습니다. MAX-THREAD 인자는 무시됩니다.")
+
+    print("THREAD_MOUNT=%d" % tcount)
 
     # 압축 파일 생성 여부
     zip_file = False
@@ -271,7 +286,7 @@ def help(error=None):
     # 연도 별로 CVE 데이터베이스 작업을 설정해주고 수행합니다.
     if year_based_task:
         for year in list_year:
-            export_filepath.append(configure_task(base, table_name, year, filename, file_path))
+            export_filepath.append(configure_task(base, table_name, year, filename, file_path, thread_count=tcount))
 
         if len(base.execption_list) != 0:
             log_filename = "err-{time}.log".format(time=datetime.now()).replace(" ", "").replace(":", ".")
@@ -287,7 +302,7 @@ def help(error=None):
             print("\n경고: {count} elements could not be inserted into the table. See {f} for more reasons.\n".format(count=len(base.execption_list), f=log_filename))
             base.execption_list.clear()
     else:
-        export_filepath.append(configure_task(base, table_name, None, filename, file_path))
+        export_filepath.append(configure_task(base, table_name, None, filename, file_path, thread_count=tcount))
 
     # 출력 파일로 만들어져야 할 것이 있습니까?
     if len(output_filetype) != 0:

database.py

@@ -197,7 +197,7 @@ def export(self, export_path_name="unspecified", export_type='sql', tablename="U
                     password_exist = "-p{pwd}".format(pwd=self.password)
                 else:
                     password_exist = ""
-                ps = subprocess.Popen("mysqldump -u {user} {pwd} -h {host} {database_name} > {output}".format(pwd=password_exist, user=self.user,host=self.host, database_name=self.database,output=filename), shell=True)
+                ps = subprocess.Popen("mysqldump --column-statistics=0 -u {user} {pwd} -h {host} {database_name} > {output}".format(pwd=password_exist, user=self.user,host=self.host, database_name=self.database,output=filename), shell=True)
                 ps.communicate()
                 if ps.returncode == 0:
                     print("Sucessfully make the export file: %s" % filename)