"""
Author : Saikiranuppu
Version: 1.0.0
Created: 2017-11-21
"""

import os
import argparse
import time
import magic
import logging
import hashlib
from shutil import copyfile,move
import json
import datetime
import MySQLdb
import logging
import sys

if not os.geteuid() == 0:
    sys.exit("\nOnly root can run this script\n")

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.FileHandler('meta.log')
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

with open('config.json','r') as f:
    config = json.load(f)

result = {}
dir_path = os.path.dirname(os.path.realpath(__file__))
default_root  = config.get('default_root',dir_path)
mysql_host    = config['mysql_host']
mysql_user    = config['mysql_user']
mysql_pass    = config['mysql_pass']
mysql_database= config['mysql_database']

def md5(fname):
    hash_md5 = hashlib.md5()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()
def sha256(fname):
    hash_sha256 = hashlib.sha256()
    with open(fname,'rb') as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_sha256.update(chunk)
    return hash_sha256.hexdigest()
def insert_db(file_meta):
    try:
        connection = MySQLdb.connect(mysql_host,mysql_user,mysql_pass,mysql_database)
        cursor = connection.cursor()
        sql = "INSERT INTO `file_meta` (`md5`,`sha256`,`mime`,`severity`,`timestamp`,`file_source`,`file_size`,`file_path`,`file_new_location`,`file_name`,`author`,`tags`) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) on DUPLICATE KEY UPDATE `timestamp`= %s"
        cursor.execute(sql,(file_meta['md5'],file_meta['sha256'],file_meta['mime'],file_meta['severity'],file_meta['timestamp'],file_meta['file_source'],file_meta['file_size'],file_meta['file_path'],file_meta['file_new_location'],file_meta['file_name'],file_meta['author'],file_meta['tags'],file_meta['timestamp']))
        connection.commit()
        logger.info('Inserting file_meta into Database {}'.format(file_meta))
    except Exception as e:
        print 'Error in DB Connection {}'.format(str(e))
    finally:
        connection.rollback()
        connection.close()


def process_files(path,source,severity,delete,author,tags):
    if os.path.isfile(path):
        pass
    for root, subFolders, files in os.walk(path):
        for file_original_name in files:
            file_meta = {}
            file_meta['file_name'] = file_original_name
            file_meta['file_path'] = os.path.join(root,file_original_name)
            file_meta['file_source'] = source
            file_meta['severity'] = severity
            file_meta['author'] = author
            file_meta['file_size'] = os.path.getsize(file_meta['file_path'])
            file_meta['md5'] = md5(file_meta['file_path'])
            file_meta['tags'] = tags
            file_meta['timestamp'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            file_meta['sha256'] = sha256(file_meta['file_path'])
            mime = magic.Magic(mime=True)
            file_meta['mime'] = mime.from_file(file_meta['file_path']).replace('/','_')
            file_meta['file_new_location'] = '{}/{}/{}/{}'.format(default_root,severity,file_meta['mime'],file_original_name)
            try:
                extracted_new_dir = os.path.join(default_root,severity,file_meta['mime'])
                if not os.path.isdir(extracted_new_dir):
                    os.makedirs(extracted_new_dir)
                if delete:
                    move(file_meta['file_path'],file_meta['file_new_location'])
                else:
                    copyfile(file_meta['file_path'],file_meta['file_new_location'])            
                
                print '[*] FileMeta : {} '.format(json.dumps(file_meta,sort_keys=True,indent=4))
            except Exception as e:
                print '[*] Exception : {}'.format(str(e))
            if config['database']:
                insert_db(file_meta)
def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-f" ,
                        "--folder",
                        required=True,
                        help="Samples Folder/File")
    parser.add_argument("-src","--source",required=True,
                        help="Source from which malware is downloaded")
    parser.add_argument("-sev","--severity",required=True,
                        choices=set(('clean','malicious')),
                        help="Benign or Malicious")
    parser.add_argument("-d","--delete",default=0,
                        help="Delete files from source folder after copying")
    parser.add_argument("-a","--author",required=True,
                        help="Author name")
    parser.add_argument("-t","--tags",default='',
                        help="Tags to classify the samples")
    args = parser.parse_args()
    source_path = args.folder
    author = args.author
    source = args.source
    severity = args.severity
    delete = args.delete
    tags = args.tags

    print '[*] Starting the Malwares Organisation Process'
    print '[*] Delete Original Files options is set to {}'.format(delete)
    process_files(source_path,source,severity,delete,author,tags)


if __name__ == '__main__':
    main()
