# This file contiains the functions used by the scraper to exfiltrate all
#   scraped data into a sql database

# Third Party Imports
import sys
import sqlite3
from sqlite3 import Error
import datetime


# Error log writer
error_log = open("./Logs/sql_log.log", "a+")


# Log errors to the file specified by error_log
def log_message(msg):
    timestamp = str(datetime.datetime.now())
    error_log.write(timestamp + ":\t" + msg + "\n")


# Close out the module cleanly
def close_database_functions():
    log_message("closing database broker cleanly")
    error_log.close()


# Establish connection with the sql database
def link_db(dbName):
    try:
        connection = sqlite3.connect(dbName)
        print("SQLITE3 VERSION:\t" + sqlite3.version + "\n\n")
        return connection
    except Error as e:
        log_message(e)

    return None


# Function that is used to help streamline sql calls and keep code neat.
def __execute_sql(database_connection, command):
    cursor = database_connection.cursor()
    if database_connection is not None:
        cursor.execute(command)
    else:
        log_message("ERROR! Unable to connect to database")


# Creates all tables and handles any other sqlite settings that will be needed
#   before we can push data in.
def initialize_database(database_connection):
    generate_cve_table_sql = """CREATE TABLE IF NOT EXISTS cve_store (
                                cve_id text PRIMARY KEY,
                                cve_page text,
                                cwe_id text,
                                known_exploits text,
                                vulnerability_classification text,
                                publish_date text,
                                update_date text,
                                score text,
                                access_gained text,
                                attack_origin text,
                                complexity text,
                                authentication_required text,
                                confidentiality text,
                                integrity text,
                                availability text);"""

    __execute_sql(database_connection, generate_cve_table_sql)


def insert_cve_into_database(connection, cve):
    sql = """INSERT INTO cve_store(cve_id,cve_page,cwe_id,
            known_exploits,vulnerability_classification,
            publish_date,update_date,score,access_gained,
            attack_origin,complexity,authentication_required,
            confidentiality,integrity,availability)
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?);"""

    cursor = connection.cursor()
    try:
        cursor.execute(sql, cve)
    except Error as e:
        print(str(e)) # set up real error table


################################################################################
# MAIN
################################################################################
def main(argv):
    print("NOTE: This file should only be directly ran for testing!")
    print("\n==== CVE-Scraper ====")
    print("==== Main.py ====\n")
    print("PYTHON VERSION:\t\t" + sys.version)

    database_name = 'cve.db'

    myDB = link_db(database_name)

    initialize_database(myDB)

    quit()


if __name__ == "__main__":
    main(sys.argv[1:])
