import gzip
from datetime import datetime
import struct, socket
import logging
import os,sys
import MySQLdb
import glob
import yaml
from sys import exc_info

appPath = os.path.abspath(os.path.dirname(os.path.realpath(__file__)) + "/../../")
sys.path.append(appPath + "/lib/pylib")

from pull_config import Configurator
configr = Configurator()

# Credentials used for the database connection
configr = Configurator()
DB = configr.get_var('db')
HOST = configr.get_var('db_host')
USERNAME = configr.get_var('db_user')
PASSWORD = configr.get_var('db_pass')


#logging set up
logger = logging.getLogger('named_malware.py')
hdlr = logging.FileHandler(appPath + '/logs/message_log')
error_hdlr = logging.FileHandler(appPath + '/logs/error_log')
formatter = logging.Formatter('%(asctime)s %(levelname)s %(name)s: %(message)s')
hdlr.setFormatter(formatter)
error_hdlr.setFormatter(formatter)
error_hdlr.setLevel(logging.ERROR)
logger.addHandler(hdlr)
logger.addHandler(error_hdlr)
logger.setLevel(logging.DEBUG)
logger.debug('args: [\''+('\', \''.join(sys.argv))+'\']')

#config vars
named_dir = configr.get_var('approot')+"app/scripts/named/"

#parse args and read import config
if len(sys.argv)< 2:
    logger.error('Config file path required')
    exit(1)

try:
    import_config=yaml.load(open(sys.argv[1],'r').read())
except:
    logger.error('Error loading config file: {0}'.format(sys.argv[1]))

whitelisted_ips = set(import_config['whitelisted_ips'])
src = import_config['named_src']
src_id = -1 #pre database lookup value
chunksize = import_config['chunksize']
archive = import_config['archive']

current_year = datetime.today().year
current_month = datetime.today().month
domains = {}
ips={}
uniq_set = set()
count = 0
conn = None
cursor = None


def connect_db():
    global conn,cursor
    logger.info("Opening database connection")
    conn = MySQLdb.connect(host=HOST,
      user=USERNAME,
      passwd=PASSWORD,
      db=DB)
    cursor = conn.cursor()

#Used add the correct year to the date because the logs do not contain the year
def convert_date(d):
    dt=datetime.strptime(d,"%b %d %H:%M:%S")
    if current_month<12 and dt.month==12:
        dt=dt.replace(year=current_year-1)
    else:
        dt=dt.replace(year=current_year)
    return dt.strftime("%Y-%m-%d %H:%M:%S")

#Check if the line is a dns resolution
def is_dns_resolution(l):
    return len(l)>8 and l[5]=='client' and l[8]=='query:'

#Checks if the record is unique
#Also add the record to the unique set if it is unique
def is_unique(record):
    global uniq_set
    if record in uniq_set:
        return False
    uniq_set.add(record)
    return True

#Extracts the date from the line
#Calls convert_date and returns the result
def get_date(l):
    date = ' '.join(l[:3])
    return convert_date(date)

#Get ip and ip numeric from the ip portion of the log line
def get_ip(ip):
    #strip off extra data
    ip = ip.split('#')[0]

    #Check if the numeric value was computed previously
    if ip in ips:
        ip_numeric=ips[ip]
    else:
        #Compute numeric value and store for quick lookup
        try:
            ip_numeric = struct.unpack('>L',socket.inet_aton(ip))[0]
        except:
            ip_numeric = -1
            logger.error('error with ip numeric for record: {0}'.format(str(l)))
            raise
        ips[ip]=ip_numeric
    return ip,ip_numeric

#Load all malicious domains from database into dictionary object for quick lookup
def load_domains():
    global domains
    query = "SELECT domain_id, domain_name from domain where domain_is_malicious > 0"
    try:
        cursor.execute(query)
    except AttributeError:
        logger.debug('no connection to db. calling connect_db')
        connect_db()
        cursor.execute(query)
    res = cursor.fetchall()
    for record in res:
        domains[record[1].lower()] = int(record[0])

#Get domain_id for a malicious domain
#Returns -1 if domain is not marked malicious in database
def get_domain_id(domain):
    global domains
    if len(domains)==0:
        load_domains()
    domain = domain.lower()
    if not domain in domains:
        return -1;
    domain_id = domains[domain]
    return domain_id


#Returns source id for the named source
def get_src_id(src):
    global src_id
    #If the source id was previously looked up return the locally saved value
    if src_id>=0:
        return src_id

    #Look up the id by source name
    query='select named_src_id from named_src where named_src_name=%s'
    try:
        cursor.execute(query,(src,))
    except:
        db_connect()
        cursor.execute(query,(src,))
    res = cursor.fetchone()

    #Source was not in database
    #Insert the source and return the new id
    if res == None:
        query = 'insert into named_src set named_src_name=%s'
        cursor.execute(query,(src,))
        conn.commit()
        src_id = int(cursor.lastrowid)

    #Source was in database
    #Save id for faster lookup
    else:
        src_id = int(res[0])

    return src_id

def proc_line(line):
    #split line on spaces
    l = line.split()
    #check if the line corresponds to a dns resolution
    if is_dns_resolution(l):
        #get id for malicious domain
        dm_id = get_domain_id(l[9])
        #If domain has id keep processing else skip record(return -1)
        if dm_id > -1:
            #Get ip and ip numeric
            ip,ip_numeric = get_ip(l[6])
            #If ip in whitelist skip record(return -1) else keep processing
            if ip not in whitelisted_ips:
                #get date
                date = get_date(l)
                #get src_id
                src_id = get_src_id(src)
                #return string formatted for import
                return ','.join(str(x) for x in [date,ip,ip_numeric,dm_id,src_id])
    #Indicate record will not be included in the import (skip record)
    return -1

#read gzipped log file line by line and write to files for archive and import
def proc_file(filepath,archive_filepath,chunk_filepath_template,chunksize):
    logger.debug('processing file: {0}'.format(filepath))
    logger.debug('chunk size: {0}'.format(chunksize))
    count = 0
    fnumber = 0
    fchunkout=None
    chunk =''
    #open gzipped log file
    fin = gzip.open(filepath,'r')
    for l in fin:
        res = proc_line(l)
        #If proc_line returned a formatted line and the line is unique
        if res != -1 and is_unique(res):
            #add line to chunk variable and inc counter
            chunk+=res+'\n'
            count+=1
            #When chunk reaches the target chunksize write to files
            if count % chunksize == 0:
                chunk_filepath = chunk_filepath_template.format(fnumber)
                write_data(chunk,archive_filepath,chunk_filepath)
                fnumber+=1
                chunk = ''

    #If there is a partial chunk at the end write it to files
    if chunk!='':
        chunk_filepath = chunk_filepath_template.format(fnumber)
        write_data(chunk,archive_filepath,chunk_filepath)
        fnumber+=1
    fin.close()
    logger.info('{0} records written to {1} chunk files'.format(count,fnumber))
    return fnumber

#writes chunks to archive(gzipped) file and chunk file
def write_data(data, archive_filepath, chunk_filepath):
    if archive:
        logger.debug('writing to archive file: {0}'.format(archive_filepath))
        archivefile = gzip.open(archive_filepath,'a')
        archivefile.write(data)
        archivefile.close()
    logger.debug('writing to chunk file: {0}'.format(chunk_filepath))
    with open(chunk_filepath,'w') as chunkfile:
        chunkfile.write(data)

#Import chunk files into database one file at a time
#chunks is a list of numbers corresponding to a chunkfile
#chunk_filepath_template is a string that will produce the
#full file path of the chunk file when given the chunk number
def import_chunks(chunks,chunk_filepath_template):
    query = "load data local infile %s into table named_resolution fields terminated by ',' lines terminated by '\n' " + \
            "(named_resolution_datetime,named_resolution_src_ip,named_resolution_src_ip_numeric,domain_id, named_src_id)"
    for i in chunks:
        logger.info('importing chunk: {0:03d}'.format(i))
        try:
            cursor.execute(query,(chunk_filepath_template.format(i),))
            conn.commit()
        except AttributeError:
            logger.debug('no connection to db. calling connect_db')
            connect_db()
            cursor.execute(query,(chunk_filepath_template.format(i),))
            conn.commit()
        except:
            logger.error('import chunks error', exc_info=True)
            raise
        logger.info('importing chunk complete.')

#Deletes chunk files
#chunks is a list of numbers corresponding to a chunkfile
#chunk_filepath_template is a string that will produce the
#full file path of the chunk file when given the chunk number
def delete_chunks(chunks,chunk_filepath_template):
    logger.info("Cleaning up chunk files.")
    for i in chunks:
        chunk_filepath=chunk_filepath_template.format(i)
        logger.debug('removing file: {0}'.format(chunk_filepath))
        os.remove(chunk_filepath)

if __name__=='__main__':
    logger.info('named.py starting')

    #Set paths for archive and chunk directories
    archive_dir=named_dir+'archive/'+src+'/'
    chunk_dir=named_dir+'chunks/'+src+'/'

    #Make sure the directories exist
    try:
         os.makedirs(archive_dir)
    except OSError:
        logger.debug('dir exists:{0}'.format(archive_dir))

    try:
         os.makedirs(chunk_dir)
    except OSError:
        logger.debug('dir exists:{0}'.format(chunk_dir))

    #for each file in the to_load directory that is from the source defined in the config
    # - Process file and create chunk files and archive file
    # - Import the chunkfiles into database
    # - Delete chunkfiles
    # - Delete original file
    files_to_load = glob.glob(named_dir+'to_load/{0}/*.{0}.log.gz'.format(src))
    if len(files_to_load) == 0:
        logger.warning('No files to load. Exiting')
        exit(0)
    logger.debug('files to load: {0}'.format(files_to_load))
    for f in files_to_load:
        uniq_set = set()
        filepath=f
        basename=os.path.basename(f).split('.log.gz')[0]
        chunk_filepath_template=chunk_dir+basename+'.{0}.csv'
        archive_filepath=archive_dir+basename+'.csv.gz'
        num_chunks = proc_file(filepath,archive_filepath,chunk_filepath_template,chunksize)
        import_chunks(xrange(num_chunks),chunk_filepath_template)
        delete_chunks(xrange(num_chunks),chunk_filepath_template)
        os.remove(f)
    conn.close()
    logger.info('named.py complete')
