import sys
import os
import numpy, pandas, pickle, sklearn.ensemble as ske
from sklearn import model_selection, tree, linear_model
from sklearn.feature_selection import SelectFromModel
import joblib
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import confusion_matrix
from sklearn.ensemble import BaggingClassifier
from sklearn.neighbors import KNeighborsClassifier
from pprint import pprint as pp
from tqdm import tqdm

def import_data(file_to_read):
    print('Importing features from {}'.format(file_to_read))
    data = pandas.read_csv(file_to_read, sep=',')
    '''
    Drop features that we don't want to learn from:
       Name: an identifier for our benefit, if an attacker isn't inept they probably
             won't actually name something "maliciousvirus.exe"
       md5:  byte sums can be useful for traditional antivirus systems, since they can
             compare files on the user's computer against known malicious sums.
       legitimate: this is the human-labeled answer for validation. If a .exe actually
             had this value in the header antivirus would not be necessary. Learning a
             model that has this value is useless.
    '''
    #the values for training/testing:
    X = data.drop(['Name','md5','legitimate'], axis=1).values
    #the answers for testing: store the answers for each file
    Y = data['legitimate'].values
    return X,Y,data

def select_features(x,y,data):
    '''
    Use a tree classifier to select the most relevent features from data.csv
    70%-30% train-test split for purposes of cross validation.
    '''
    feature_select = ske.ExtraTreesClassifier().fit(x,y)
    model = SelectFromModel(feature_select, prefit=True)
    x_new = model.transform(x)
    nb_features = x_new.shape[1]
    x_train, x_test, y_train, y_test = model_selection.train_test_split(x_new, y, test_size=0.3)
    features = []
    # print('{} features were selected as being important:'.format(nb_features))
    indices = numpy.argsort(feature_select.feature_importances_)[::-1][:nb_features]
    col_width = len(max(data.columns[2+indices[f]] for f in range(nb_features))) + 5

    imp_features = []
    for f in range(nb_features):
        number = f+1
        feature_name = ''.join(data.columns[2+indices[f]].ljust(col_width))
        feature_importance = feature_select.feature_importances_[indices[f]]
        # print('{}.\t{} {}'.format(number, feature_name, (feature_importance * 100)))
        imp_features.append([str(feature_name).strip(), feature_importance * 100])

    for f in sorted(numpy.argsort(feature_select.feature_importances_)[::-1][:nb_features]):
        features.append(data.columns[2+f])
    # pp(imp_features)
    # return x_train, x_test, y_train, y_test, features
    return imp_features

if __name__ == "__main__":
    features = {}
    x = import_data('cleaned_data.csv')
    columns = [
            "Name",
            "md5",
            "Machine",
            "SizeOfOptionalHeader",
            "Characteristics",
            "MajorLinkerVersion",
            "MinorLinkerVersion",
            "SizeOfCode",
            "SizeOfInitializedData",
            "SizeOfUninitializedData",
            "AddressOfEntryPoint",
            "BaseOfCode",
            "BaseOfData",
            "ImageBase",
            "SectionAlignment",
            "FileAlignment",
            "MajorOperatingSystemVersion",
            "MinorOperatingSystemVersion",
            "MajorImageVersion",
            "MinorImageVersion",
            "MajorSubsystemVersion",
            "MinorSubsystemVersion",
            "SizeOfImage",
            "SizeOfHeaders",
            "CheckSum",
            "Subsystem",
            "DllCharacteristics",
            "SizeOfStackReserve",
            "SizeOfStackCommit",
            "SizeOfHeapReserve",
            "SizeOfHeapCommit",
            "LoaderFlags",
            "NumberOfRvaAndSizes",
            "SectionsNb",
            "SectionsMeanEntropy",
            "SectionsMinEntropy",
            "SectionsMaxEntropy",
            "SectionsMeanRawsize",
            "SectionsMinRawsize",
            "SectionMaxRawsize",
            "SectionsMeanVirtualsize",
            "SectionsMinVirtualsize",
            "SectionMaxVirtualsize",
            "ImportsNbDLL",
            "ImportsNb",
            "ImportsNbOrdinal",
            "ExportNb",
            "ResourcesNb",
            "ResourcesMeanEntropy",
            "ResourcesMinEntropy",
            "ResourcesMaxEntropy",
            "ResourcesMeanSize",
            "ResourcesMinSize",
            "ResourcesMaxSize",
            "LoadConfigurationSize",
            "VersionInformationSize"
        ]
    for c in columns:
        features[c] = 0
    arr = []
    for i in tqdm(range(100)):
        y = select_features(x[0], x[1], x[2])
        arr.append(len(y))
        for a in y:
            features[a[0]] += a[1]
    pp(features)