import glob
import gzip
import os
import pymysql
import logging
import subprocess

# path to gz directory
GZ_DIR = '/Users/mani/Desktop/mysql/csv'
# Database Infomation
DB_HOST='10.1.1.38'
DB_USER='diaguser'
DB_PASS='Password'
DB_NAME='PADIAGDB'
LOGFILE="exception.log"

def csv_reader(file, header=False):
    import csv
    with open(file, "r") as f:
        reader = csv.reader(f)

        if header:
            next(reader)

        for row in reader:
            yield row

# function to connect sql
def import_sql(filename, dbHostName, dbUser, dbPassword, databaseName):
    #Prepare SQL query to INSERT a record into the database.

    # Open database connection
    db = pymysql.connect(host=dbHostName,
                         user=dbUser,
                         password=dbPassword,
                         db=databaseName,
                         charset='utf8')

    for row in csv_reader(filename, False):
        # prepare a cursor object using cursor() method
        with db.cursor() as cursor:
            if row[3] == "THREAT" and row[4] == "url":
                sql = ("INSERT INTO PADIAGDB.url ("
                       "Domain,Receive_Time,Serial,Type,Threat_Content_Type,"
                       "Config_Version,Generate_Time,Source_address,Destination_address,"
                       "NAT_Source_IP,NAT_Destination_IP,Rule,Source_User,"
                       "Destination_User,Application,Virtual_System,Source_Zone,"
                       "Destination_Zone,Inbound_Interface,Outbound_Interface,Log_Action,"
                       "Time_Logged,Session_ID,Repeat_Count,Source_Port,Destination_Port,"
                       "NAT_Source_Port,NAT_Destination_Port,Flags,IP_Protocol,Action,"
                       "URL_Filename,Threat_Content_Name,Category,Severity,Direction,"
                       "Sequence_Number,Action_Flags,Source_Country,Destination_Country,"
                       "cpadding,contenttype,pcap_id,filedigest,cloud,url_idx,user_agent,"
                       "filetype,xff,referer,sender,subject,recipient,reportid,"
                       "dg_hier_level_1,dg_hier_level_2,dg_hier_level_3,dg_hier_level_4,"
                       "Virtual_System_Name,Device_Name,file_url )"
                       ""
                       "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s"
                       "); ")

            elif row[3] == "THREAT":
                sql = ("INSERT INTO PADIAGDB.threat ("
                       "Domain,Receive_Time,Serial,Type,Threat_Content_Type,"
                       "Config_Version,Generate_Time,Source_address,Destination_address,"
                       "NAT_Source_IP,NAT_Destination_IP,Rule,Source_User,"
                       "Destination_User,Application,Virtual_System,Source_Zone,"
                       "Destination_Zone,Inbound_Interface,Outbound_Interface,Log_Action,"
                       "Time_Logged,Session_ID,Repeat_Count,Source_Port,Destination_Port,"
                       "NAT_Source_Port,NAT_Destination_Port,Flags,IP_Protocol,Action,"
                       "URL_Filename,Threat_Content_Name,Category,Severity,Direction,"
                       "Sequence_Number,Action_Flags,Source_Country,Destination_Country,"
                       "cpadding,contenttype,pcap_id,filedigest,cloud,url_idx,user_agent,"
                       "filetype,xff,referer,sender,subject,recipient,reportid,"
                       "dg_hier_level_1,dg_hier_level_2,dg_hier_level_3,dg_hier_level_4,"
                       "Virtual_System_Name,Device_Name,file_url )"
                       ""
                       "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s"
                       "); ")

            elif row[3] == "TRAFFIC":
                sql = ("INSERT INTO PADIAGDB.traffic ("
                       "Domain,Receive_Time,Serial,Type,Threat_Content_Type,Config_Version,"
                       "Generate_Time,Source_address,Destination_address,NAT_Source_IP,"
                       "NAT_Destination_IP,Rule,Source_User,Destination_User,Application,"
                       "Virtual_System,Source_Zone,Destination_Zone,Inbound_Interface,"
                       "Outbound_Interface,Log_Action,Time_Logged,Session_ID,Repeat_Count,"
                       "Source_Port,Destination_Port,NAT_Source_Port,NAT_Destination_Port,"
                       "Flags,IP_Protocol,Action,Bytes,Bytes_Sent,Bytes_Received,Packets,"
                       "Start_Time,Elapsed_Time,Category,Padding,Sequence_Number,Action_Flags,"
                       "Source_Country,Destination_Country,cpadding,pkts_sent,pkts_received,"
                       "session_end_reason,dg_hier_level_1,dg_hier_level_2,dg_hier_level_3,"
                       "dg_hier_level_4,Virtual_System_Name,Device_Name,action_source )"
                       ""
                       "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s"
                       "); ")
            elif row[3] == "SYSTEM":
                sql = ("INSERT INTO PADIAGDB.system ("
                       "Domain,Receive_Time,Serial,Type,Threat_Content_Type,Config_Version,"
                       "Generate_Time,Virtual_System,Event_ID,Object,fmt,id,module,Severity,"
                       "Description,Sequence_Number,Action_Flags,dg_hier_level_1,"
                       "dg_hier_level_2,dg_hier_level_3,dg_hier_level_4,Virtual_System_Name,"
                       "Device_Name )"
                       ""
                       "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,"
                       "%s,%s,%s );")
            else:
                continue

            try:
                cursor.execute('SET foreign_key_checks = 0')
                # Execute the SQL command
                r = cursor.execute(sql, row)

                #Commit your changes in the database
                cursor.execute('SET foreign_key_checks = 1')
                db.commit()


            except Exception as e:
                logging.exception(e)
                db.rollback()

    # disconnect from server
    db.close()

#####################################################################
# Main
#####################################################################
logging.basicConfig(level=logging.DEBUG, filename=LOGFILE)
# # List of all gz files
# gz_files = (gz for gz in glob.glob(os.path.join(GZ_DIR, '*.gz')))
#
# # Step 4: in loop, back to step 1 automatically for another gz files
# for gz_file in gz_files:
#     # Step 1
#     with gzip.open(gz_file, 'rb') as in_file:
#         s = in_file.read()
#
#         # store the uncompressed data
#         sql_file = gz_file[:-3]
#
#         #Convert to csv
#         sql_file= sql_file[:-4] + '.csv'
#         with open(sql_file, 'wb') as out_file:
#             out_file.write(s)
#
#     # Step 2: import to sql
#     # Fill value for dbHostname, dbUser, dbPassword, databaseName
#     import_sql(out_file,DB_HOST,DB_USER,DB_PASS,DB_NAME)
#
#     # Step 3: remove uncompresed file
#     os.remove(sql_file)
#
#     # Step 4: in loop, back to step 1 automatically for another gz files

#####################################################################
# Main
#####################################################################

# List of all gz files
gz_files = (gz for gz in glob.glob(os.path.join(GZ_DIR, '*.gz')))

# Loop through all gz files
for gz_file in gz_files:
    cmd = ['gz', 'a', 'Test.7z', 'Test', '-mx9']
    subprocess.call(cmd)
    # Step 1
    with gzip.open(gz_file, 'rb') as in_file:
        s = in_file.read()

    # store the uncompressed data
    sql_file = gz_file[:-3]

    # Convert to csv
    sql_file = sql_file[:-4] + '.csv'
    with open(sql_file, 'wb') as out_file:
        out_file.write(s)

    # Step 2: import to sql
    import_sql(out_file, DB_HOST, DB_USER, DB_PASS, DB_NAME)

    # Step 3: remove uncompresed file
    os.remove(sql_file)

    # Step 4: in