cvemitre.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import csv
import re
import codecs
import xml.etree.ElementTree as ET
import unicodedata
from database import *
from cve import CVEDatabase

execption_list = {}

class Cvemitre(CVEDatabase):
    def __init__(self, mysql_base):
        super().__init__(mysql_base)
    
    def default_table_column(self):
        # To check value using increment number.
        # a1 = TableColumn("_number", "INT PRIMARY KEY AUTO_INCREMENT")

        a2 = TableColumn("cve_year", "VARCHAR(5)")
        a3 = TableColumn("name", "VARCHAR(24)")
        a4 = TableColumn("status", "VARCHAR(12)")
        a5 = TableColumn("description", "TEXT")
        a6 = TableColumn("reference", "TEXT")
        a9 = TableColumn("phase", "VARCHAR(32)")
        a7 = TableColumn("vote", "TEXT")
        a8 = TableColumn("comment", "TEXT")
        return tuple(vars().values())

    def load(self, obj):
        self.connect()
        if isinstance(obj, dict):
            year = obj.get('year')
            raw_type = obj.get('specific_type')
        else:
            raise TypeError("The variable 'obj' must be dict, not {t}".format(t=type(obj).__str__()))
        
        url = "https://cve.mitre.org/data/downloads/"

        if raw_type is None:
            raw_type = "csv"
        raw_file = "allitems.{ft}".format(ft=raw_type)
        is_exist = False
        
        if os.path.exists(raw_file):
            if raw_type == "xml":
                if util.vaildxml(raw_file):
                    print("The resource file already exist, continue without the download")
                    is_exist = True
                else:
                    os.remove(raw_file)
            elif raw_type == "csv":
                if util.vaildcsv(raw_file):
                    print("The resource file already exist, continue without the download")
                    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 ...")
            util.downloadfrom(target_url)
        
        print("Converting to UTF-8 strings ...")

        BLOCKSIZE = 2057192
        if os.path.exists(raw_file + ".conv"):
            os.remove(raw_file + '.conv')
            print("Found the temp file, Removed it to read the file to safe.")

        with codecs.open(raw_file, "r", "iso-8859-1") as sf:
            with codecs.open(raw_file + ".conv", "w", "utf-8") as tf:
                while True:
                    contents = sf.read(BLOCKSIZE)
                    if not contents:
                        break
                    tf.write(contents)

        collection = []
        if raw_type == "xml":
            readTo = ET.parse(raw_file).getroot()
            raise NotImplementedError("XML Type is not implemented")
        else:
            with open(raw_file + ".conv", 'r', newline='', encoding='utf-8') as fs:
                    readTo = csv.reader(fs)
                    print("Analyzing CSV data ...")
                    for line in readTo:
                        try:
                            name = line[0]
                            if not name.startswith("CVE-"):
                                continue
                            status = line[1]
                            description = line[2]
                            reference = line[3]
                            phase = line[4]
                            vote = line[5]
                            comment = line[6]
                            collection.append((name,status,description,reference,phase,vote,comment))
                        except UnicodeDecodeError as e:
                            print("오류: 해당 CSV Element를 객체화할 수 없었습니다 -> {cve} | Reason: {reason}".format(cve=name, reason=e)) 
        if os.path.exists(raw_file + ".conv"):
            print("Removing temp file ...")
            os.remove(raw_file + '.conv')

        return raw_file, collection

    def table_insert(self, tablename, collection):
        print("Inserting data ...")
        for line in collection:
            line = list(line)
            name = line[0]
            p = re.compile(r'CVE-(\d{4})-.+$')
            cve_year = str(p.match(name).group(1))
            status = line[1]
            description = line[2]
            if len(line[3]) != 0:
                try:
                    reference = unicodedata.normalize('NFKD', line[3]).encode('utf-8')
                except UnicodeDecodeError:
                    print("Warning: {N} - There's does not support UTF-8, some values may be lost".format(N=name))
                    reference = ''
                    for p in line[3]:
                        try:
                            bytes(p, 'utf-8')
                        except UnicodeEncodeError:
                            p = '?'
                        reference = reference + p
            phase = line[4]
            vote = line[5]
            comment = line[6]

            # Initializes all temporary variables to read only meaningful variables in the function.
            ex = None
            query_argument = None
            query_format = None
            query_argument = None
            query_header = None
            variable_name = None
            detect = None
            p = 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 += ")"
            self.send_query("INSERT INTO {table_name} {header} VALUES {format}".format(table_name=tablename, header=query_header, format=query_format), tuple(query_argument), False, execption_list)
            if len(execption_list) != 0:
                current_key = list(execption_list.keys())
                for ex in current_key:
                    print("실패: Column에 값을 넣을 수 없습니다. -> {cve} | Reason: {r}".format(cve=name, r=ex))
                execption_list.clear()

    def insert(self, obj):
        raise NotImplementedError("Not 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, charset)