import numpy as np
import pandas as pd
import argparse
import joblib
import glob2
import hashlib
import os
import pefile
import entropy
import string
import argparse
from sklearn.feature_extraction.text import CountVectorizer
import csv
import time

field_names = ['name', 'strings', 'sha256', 'imports', 'size', 'signed', 'good_file', 'entropy']


def get_file_size(FilePath):
    return os.path.getsize(FilePath)


def get_file_signatures(FilePath):

    try:
        pe = pefile.PE(FilePath)

        address = pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_SECURITY']].VirtualAddress
        size = pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_SECURITY']].Size
        if address == 0:
            return 0
        return 1
    except:
	    return -1

def get_sha1(FilePath):
    return hashlib.sha1(FilePath).hexdigest()

def get_md5_of_file(FilePath):
    return hashlib.md5(FilePath).hexdigest()


def get_sha256_of_file(FilePath):
    return hashlib.sha256(FilePath).hexdigest()


def get_filename_from_path(FilePath):
    FileSplitPath = FilePath.split("\\")
    return FileSplitPath[-1]


def get_file_path(FilePath):
    FileSplitPath = FilePath.split("\\")
    seq = "\\"
    return seq.join(FileSplitPath[0:len(FileSplitPath) - 1])


def get_import_table_entries(FilePath):
    imports = []

    try:
        pe = pefile.PE(FilePath)

        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            # print entry.dll
            for imp in entry.imports:
                imports.append(imp.name)
    except:
        # print "Error: ", filepath
        imports.append("Error")
    return imports


def get_entropy_of_file(FilePath):
    file_read = open(FilePath, 'r')
    file_data = file_read.read()

    entropy_file = entropy.shannon_entropy(file_data)
    file_read.close()
    return entropy_file


def get_type_of_file(FilePath):
    magic_type = magic.from_file(FilePath)

    if 'DLL' in magic_type:
        return 0
    else:
        return 1


def is_pefile(FilePath):
    is_PE = 0
    file_handle = open(FilePath, "rb")

    if file_handle.read(2) == "MZ":
        is_PE = 1
    file_handle.close()

    return is_PE

def check_whitelist(sha1, whitelist):
    return 1

def get_strings(filename, min_length):
    file_data = open(filename, 'rb')

    strings_list = []

    current = ''
    for character in file_data.read():
        if character in string.printable:
            current += character
        else:
            if len(current) > min_length:
                strings_list.append(current)

            current = ''

    return strings_list

def main():
    parser = argparse.ArgumentParser(description="File Scanner for Blue Sentry")
    parser.add_argument("--imports_encoder", help="Input Encoder Pickle File", required=True)
    parser.add_argument("--strings_encoder", help="Input Encoder Pickle File", required=True)
    parser.add_argument("--imports_classifier", help="Import Classifier Pickle File", required=True)
    parser.add_argument("--strings_classifier", help="Import Classifier Pickle File", required=True)
    parser.add_argument("--stage2_classifier", help="Stage 2 Classifier Pickle File", required=True)
    parser.add_argument("--whitelist", help="Whitelist of files that are known good", required=True)
    parser.add_argument("--directory", help="Directory you wish to scan")
    parser.add_argument("--file", help="The file you wish to scan")
    args = parser.parse_args()

    print args
    csv_file = open("testfile.csv", 'a+')

    writer = csv.DictWriter(csv_file, field_names)
    writer.writeheader()
    files_to_scan = []

    if args.file is not None:
        files_to_scan = glob2.glob(args.file)
    elif args.directory is not None:
        files_to_scan = glob2.glob(args.directory)
        print files_to_scan
    else:
        print "Need to input a file or directory"
        return

    encoder_imports = joblib.load(args.imports_encoder)
    encoder_strings = joblib.load(args.strings_encoder)
    imports_clf = joblib.load(args.imports_classifier)
    final_clf = joblib.load(args.stage2_classifier)

    for filename in files_to_scan:

        print filename
        size = get_file_size(filename)
        imports = get_import_table_entries(filename)
        sha256 = get_sha256_of_file(filename)
        signed = get_file_signatures(filename)
        entropy = get_entropy_of_file(filename)
        strings = get_strings(filename, 8)
        sha1 = get_sha1(filename)
        #file_type = get_type_of_file(filename)
        good_file = 0
        if is_pefile(filename):
            writer.writerow(
                {'name': filename, 'strings': strings, 'sha256': sha256, 'imports': imports, 'size': size,
                 'signed': signed, 'good_file': good_file, 'entropy': entropy})

    csv_file.close()
    data = pd.read_csv("testfile.csv")
    data.columns = field_names
    #print data
    data_imports = data['imports']
    data_imports_encoded = encoder_imports.transform(data_imports)

    pred_prob = imports_clf.predict_proba(data_imports_encoded)

    pred_prob_rounded = []
    for x in range(0, data_imports_encoded.shape[0]):
        pred_prob_rounded.append(round(pred_prob[x][1], 8))

    data['imports_processed'] = pred_prob_rounded

    clf_strings = joblib.load(args.strings_classifier)
    data_strings_encoded = encoder_strings.transform(data['strings'])
    pred_prob_strings = clf_strings.predict_proba(data_strings_encoded)
    pred_prob_rounded_strings = []
    for x in range(0, data_strings_encoded.shape[0]):
        pred_prob_rounded_strings.append(round(pred_prob_strings[x][1], 8))

    data['strings_processed'] = pred_prob_rounded_strings

    data = data[['name', 'sha256', 'size', 'signed', 'good_file', 'entropy', 'imports_processed', 'strings_processed']]

    data_names = data['name']
    data = data.drop(['sha256'], axis=1)
    data = data.drop(['size'], axis=1)
    print data
    data = data.drop('name', axis=1)
    data = data.drop(['good_file'], axis=1)
    #data = data.drop(['file_type'], axis=1)

    predictions = final_clf.predict(data)

    index = 0
    for name in data_names:
        print "Index: " + str(index)
        if predictions[index] == 1:
            print "%s was marked as good" % name
        else:
            print "%s was marked as bad" % name
        index += 1
    os.remove("testfile.csv")

if __name__ == "__main__":
    main()
