import re
from csv import reader
from colorama import init, Fore

# Insert the actual exploits in searchsploit in the database
def update_database(exploit_database, mycursor):
    print(Fore.BLUE + "Updating database...")
    
    # Read the CSV to get the basic information
    with open('/usr/share/exploitdb/files_exploits.csv','r') as read_obj:
        # Read the CSV and skip the first row (headers)
        csv_reader = reader(read_obj)
        next(csv_reader)

        # Insert each row in the table 
        for row in csv_reader:
            query = """INSERT IGNORE INTO Exploits (ID, File, Description, Date, Author, Type, Platform, Port, SellerLink, SoftwareLink, Version, Tested, CVE)
                       VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                       """
            values = (row[0],row[1],row[2],row[3],row[4],row[5],row[6],row[7])

            # To get more information about the exploit 
            values = search_content(row[1], values)

            mycursor.execute(query, values)
            
    exploit_database.commit()
    print(Fore.GREEN + "Database update")

# Search the exploit content to find more information about it.
def search_content(exploit_path, values):
    # Add the root path to the exploit path
    exploit_path = "/usr/share/exploitdb/" + exploit_path

    # Specific variables to look for within each exploit
    seller_link = ""
    software_link = ""
    version = ""
    tested = ""
    CVE = ""

    # Booleans to control the search
    isEmptyVendor = True
    isEmptySoftware = True
    isEmptyVersion = True
    isEmptyTested = True
    isEmptyCVE = True

    # Open the file to read its content
    with open(exploit_path, 'r') as exploit:
        # Get all the lines of the file
        data = exploit.read().splitlines()
        
        # Iterate through them to find the key words and, after cleaning it, store them
        for line in data:             
            # Search and check the vendor link
            if isEmptyVendor:
                if re.search('[Vv]endor [Hh]omepage', line): 
                    if (re.split('[Vv]endor [Hh]omepage', line)[1].strip().startswith(':')):
                        seller_link = clean_characters(line,'[Vv]endor [Hh]omepage',':')
                    elif (re.split('[Vv]endor [Hh]omepage', line)[1].strip().startswith('-')):
                        seller_link = clean_characters(line,'[Vv]endor [Hh]omepage','-')
                    elif (re.split('[Vv]endor [Hh]omepage', line)[1].startswith(' ')):
                        seller_link = clean_white(line,'[Vv]endor [Hh]omepage')
                    isEmptyVendor = False
                elif re.search('[Vv]endor', line):
                    if (re.split('[Vv]endor', line)[1].strip().startswith(':')):
                        seller_link = clean_characters(line,'[Vv]endor',':')
                    elif (re.split('[Vv]endor', line)[1].strip().startswith('-')):
                        seller_link = clean_characters(line,'[Vv]endor','-')
                    elif (re.split('[Vv]endor', line)[1].startswith(' ')):
                        seller_link = clean_white(line,'[Vv]endor')
                    isEmptyVendor = False
                    
            # Search and check the software link
            if isEmptySoftware:
                if re.search('[Ss]oftware [Ll]ink', line):
                    if (re.split('[Ss]oftware [Ll]ink', line)[1].strip().startswith(':')):
                        software_link = clean_characters(line,'[Ss]oftware [Ll]ink',':')
                    elif (re.split('[Ss]oftware [Ll]ink', line)[1].strip().startswith('-')):
                        software_link = clean_characters(line,'[Ss]oftware [Ll]ink','-')
                    elif (re.split('[Ss]oftware [Ll]ink', line)[1].startswith(' ')):
                        software_link = clean_white(line,'[Ss]oftware [Ll]ink')
                    isEmptySoftware = False
                elif re.search('[Pp]roduct [Ww]eb [Pp]age', line):
                    if (re.split('[Pp]roduct [Ww]eb [Pp]age', line)[1].strip().startswith(':')):
                        software_link = clean_characters(line,'[Pp]roduct [Ww]eb [Pp]age',':')
                    elif (re.split('[Pp]roduct [Ww]eb [Pp]age', line)[1].strip().startswith('-')):
                        software_link = clean_characters(line,'[Pp]roduct [Ww]eb [Pp]age','-')
                    elif (re.split('[Pp]roduct [Ww]eb [Pp]age', line)[1].startswith(' ')):
                        software_link = clean_white(line,'[Pp]roduct [Ww]eb [Pp]age')
                    isEmptySoftware = False

            # Search and check the affected version
            if isEmptyVersion:
                if re.search('[Vv]ersion', line):
                    if (re.split('[Vv]ersion', line)[1].strip().startswith(':')):
                        version = clean_characters(line,'[Vv]ersion',':')
                    elif (re.split('[Vv]ersion', line)[1].strip().startswith('-')):
                        version = clean_characters(line,'[Vv]ersion','-')
                    elif (re.split('[Vv]ersion', line)[1].startswith(' ')):
                        version = clean_white(line,'[Vv]ersion')
                    isEmptyVersion = False

            # Search and check where it has been tested
            if isEmptyTested:
                if re.search('[Tt]ested [Oo]n', line):
                    if (re.split('[Tt]ested [Oo]n', line)[1].strip().startswith(':')):
                        tested = clean_characters(line,'[Tt]ested [Oo]n',':')
                    elif (re.split('[Tt]ested [Oo]n', line)[1].strip().startswith('-')):
                        tested = clean_characters(line,'[Tt]ested [Oo]n','-')
                    elif (re.split('[Tt]ested [Oo]n', line)[1].startswith(' ')):
                        tested = clean_white(line,'[Tt]ested [Oo]n')
                    isEmptyTested = False

            # Search and check the CVE
            if isEmptyCVE:
                if line.__contains__('CVE ID'):
                    if (line.partition('CVE ID')[2].strip().startswith(':')):
                        CVE = clean_characters(line,'CVE ID',':')
                    elif (line.partition('CVE ID')[2].strip().startswith('-')):
                        CVE = clean_characters(line,'CVE ID','-')
                    elif (line.partition('CVE ID')[2].startswith(' ')):
                        CVE = clean_white(line,'CVE ID')
                    isEmptyCVE = False
                elif line.__contains__('CVE'):
                    if (line.partition('CVE')[2].strip().startswith(':')):
                        CVE = clean_characters(line,'CVE',':')
                    elif (line.partition('CVE')[2].strip().startswith('-')):
                        CVE = clean_characters(line,'CVE','-')
                    elif (line.partition('CVE')[2].startswith(' ')):
                        CVE = clean_white(line,'CVE')
                    isEmptyCVE = False

        # Add the new values to the values tuple
        values = values + (seller_link,)
        values = values + (software_link,)
        values = values + (version,)
        values = values + (tested,)
        values = values + (CVE,)

    return values

# Clean with characters
def clean_characters(line, word, character):
    if (word == 'CVE' or word == 'CVE ID'):
        return line.partition(word)[2].split(character,1)[1].translate({ord(i): None for i in '[]'}).strip()
    else: 
        return re.split(word, line)[1].split(character,1)[1].translate({ord(i): None for i in '[]'}).strip()

# Clean with white space
def clean_white(line, word):
    if (word == 'CVE' or word == 'CVE ID'):
        return line.partition(word)[2].translate({ord(i): None for i in '[]'}).strip()
    else:
        return re.split(word, line)[1].translate({ord(i): None for i in '[]'}).strip()
