testSalience.py

import pickle, gzip, glob, sys, keras, os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # gets rid of AVX message
import random as rn
import numpy as np
import tensorflow as tf
os.environ['PYTHONHASHSEED'] = '0'
np.random.seed(37)
rn.seed(1254)
tf.set_random_seed(89)
from keras import optimizers
from keras import backend as K
from keras.models import load_model
from keras.layers import *
from keras.models import Sequential
from keras.losses import weighted_categorical_crossentropy
from keras.callbacks import CSVLogger, ModelCheckpoint
from keras.regularizers import *
from keras.utils.generic_utils import get_custom_objects
from keras.layers.advanced_activations import LeakyReLU, ELU
sys.path.insert(0, r'.\libraries')
from kerasLayers import *
from kerasExtras import *
elu = ELU(1)
elu.__name__ = "ELU"

import time
from keras import Model
import matplotlib.pyplot as plt
from scipy.ndimage.filters import gaussian_filter
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import seaborn as sns

from blobifier import Blobifier
from sequential import SequentialClusterer

input_length = None


# binary standard
line_length = 25
num_samples = 7500
num_samples_valid = 1600
train_name = r".\libraries\datasets\kagglewindows\windows_exe_dll_kaggle_nopad_pooled.pklz"
valid_name = r".\libraries\datasets\kagglewindows\windows_exe_dll_kaggle_validation_nopad_pooled.pklz"


steps_per_epoch = num_samples/batch_size
valid_steps = num_samples_valid/batch_size  # should be this


model = load_model(r".\networks\dist binary final nets\pruned from 5\KaggleConv-22.hdf5",
                   custom_objects={'DecayingConvLSTM2D':MinConvRNN,
                                   'window_size': window_size ,
                                   'ELU': elu,
                                   }
                   )

# mode = "kaggle"
mode = "binary"
if mode == "binary":
    lossFunc = 'binary_crossentropy'
    generatorFunc = loadDataGeneratorBinary
elif mode == "kaggle":
    lossFunc = 'categorical_crossentropy'
    generatorFunc = loadDataGenerator


        
    
    

def compileModel(model):
    optimizer = "rmsprop"
    
    model.compile(optimizer=optimizer,
                  loss=lossFunc,
                  metrics=['accuracy'])



def salienceTest(model, confMatrixFile):
    if mode == "binary" or mode == "kaggle":
        train_gen = generatorFunc(train_name, num_samples)
        valid_gen = generatorFunc(valid_name, num_samples_valid)
    else: 
        raise ValueError("mode must be kaggle or binary")
    
    compileModel(model)
    print(model.summary())

    answers = []
    preds = []
    total_correct = 0
    total_run = 0
    correctCountClasses = [0]*9
    incorrectCountClasses = [0]*9
    times = []
    confMatrix = np.zeros((9,9))
    
    # remainingAmounts = []
    totalBytes = 0
    totalNNZ = 0

    grad = K.gradients(model.layers[-1].input, model.layers[1].output)[0]
    sess = K.get_session()
    
    mal_x = []
    ben_x = []
    blobifyX = []
    blobifyY = []
    scatterLabels = []
    new_mal_x = []
    new_ben_x = []
    


    blobifier = Blobifier()

    
    
    for x in range(0, int(valid_steps)):
        train_x, train_y = next(valid_gen)
    # for x in range(0, int(steps_per_epoch)):
        # train_x, train_y = next(train_gen)
        length = train_x.shape[1]
        answer = train_y.tolist()[0]
        answers.append(answer)
        
        
        out = sess.run(grad, feed_dict={model.input: train_x})
        out = out[0] # comes in array of length 1
        out = out.reshape((int(out.shape[0]*out.shape[1]), out.shape[2], out.shape[3]))
        

        # filter input based on gradient value
        new_train_x = np.copy(train_x[0])
        new_train_x[np.abs(np.max(out, axis=2)) < 1e-16] = 0   # works great
        
        
        '''
        # different filtering attempts

        # new_train_x[np.abs(np.max(out, axis=2)) > 1e-16] = 0   # works TERRIBLY for proof!
        # new_train_x[np.abs(np.max(out, axis=2)) < 1e-4] = 0   # works fine
        # new_train_x[np.max(np.abs(out), axis=2) < 1e-4] = 0   # works fine
        
        # did not work
        # m = np.max(np.abs(out), axis=2)
        # new_train_x[m < (np.mean(m) - .5*np.std(m))] = 0
        
        # attempt - this doesn't work
        # blurred = gaussian_filter(new_train_x, sigma=1)
        # new_train_x_blurred = np.copy(train_x[0])
        # new_train_x_blurred[np.abs(blurred) < 15] = 0 
        # new_train_x = new_train_x_blurred
        '''
        
        # keep filtered w/o blobs removed
        not_removed = new_train_x

        # remove low gradient areas
        if np.sum(new_train_x) > 0:
            new_train_x = blobifier.blobify(new_train_x, int(answer))
        
        
        
        
        
        # Plot images
        # fig = plt.figure()
        
        # show different saliency maps
        # ax1 = fig.add_subplot(1,4,1)
        # ax1.imshow(train_x[0], cmap='gray')
        # plt.axis('off')
        # ax2 = fig.add_subplot(1,4, 2)
        # ax2.imshow(np.mean(out, axis=2), cmap='gray')
        # plt.axis('off')
        # ax3 = fig.add_subplot(1,4, 3)
        # ax3.imshow(np.max(out, axis=2), cmap='gray')
        # plt.axis('off')
        # ax4 = fig.add_subplot(1,4, 4)
        # ax4.imshow(np.min(out, axis=2), cmap='gray')
        # plt.axis('off')
        
        # filtering example
        # ax1 = fig.add_subplot(1,4,1)
        # ax1.imshow(train_x[0], cmap='gray')
        # plt.axis('off')
        # ax3 = fig.add_subplot(1,4, 2)
        # ax3.imshow(np.max(out, axis=2), cmap='gray')
        # plt.axis('off')
        # ax1 = fig.add_subplot(1,4,3)
        # ax1.imshow(not_removed, cmap='gray')
        # plt.axis('off')
        # ax1 = fig.add_subplot(1,4,4)
        # ax1.imshow(new_train_x, cmap='gray')
        # plt.axis('off')
        
        # plt.show()
        
        
        
        
        # make prediction on filtered version
        start = time.time()
        pred = model.predict(np.asarray([new_train_x.tolist()])).tolist()[0]
        # pred = model.predict(train_x).tolist()[0]
        amt = time.time() - start
        times.append(amt)
        preds.append(pred)
        
        
        totalNNZ += np.count_nonzero(new_train_x) 
        totalBytes += train_x.size
        
        if int(answer):
            mal_x.append(train_x)
            new_mal_x.append(new_train_x)
        else:
            ben_x.append(train_x)
            new_ben_x.append(new_train_x)
        
        
        if mode == "binary":
        
            pred = round(pred[0])  
            
            
            ansClass = int(answer)
            
            if pred == ansClass:
                total_correct += 1
                correctCountClasses[ansClass] += 1
            else:
                incorrectCountClasses[ansClass] += 1
            
            # print(pred == ansClass, np.count_nonzero(new_train_x) / new_train_x.size) 
            confMatrix[ansClass][pred] += 1
            total_run += 1
            if x % 50 == 0: print("interval", x, "correct so far", total_correct, "% of total bytes left", totalNNZ/totalBytes)
        
        elif mode == "kaggle":
            # ansClass = numpy.argmax(train_y,1)[0]
            ansClass = answer.index(max(answer))
            predClass = pred.index(max(pred)) 
            
            if predClass == ansClass:
                total_correct += 1
                correctCountClasses[ansClass] += 1
            else:
                incorrectCountClasses[ansClass] += 1
                
            confMatrix[ansClass][predClass] += 1
            total_run += 1
            if x % 500 == 0: print("interval", x, "correct so far", total_correct)

    print("correct:", total_correct, "out of", total_run)
    print("correct per class  :", correctCountClasses)
    print("incorrect per class:", incorrectCountClasses)
    print("mean time to predict", np.mean(times))
    

    print(confMatrix)
    if confMatrixFile:
        np.savetxt(confMatrixFile, confMatrix, delimiter=",")



    # plot length of file vs amount removed, per class
    # plt.scatter(blobifier.blobifyX, blobifier.blobifyY, c=blobifier.blobifyC)
    # plt.show()
    
    # blobs are separated into lists - malware blobs and benign blobs
    # can use them however you want
    print("number of malware blobs:", len(blobifier.malwareBlobs))  #6019, 5982, 5957
    print("number of benign blobs:", len(blobifier.benignBlobs))  #10512, 10580
    
    # show individual blobs
    for blob in blobifier.malwareBlobs:
        fig = plt.figure()
        ax1 = fig.add_subplot(1,4,1)
        ax1.imshow(blob, cmap='gray')
        plt.axis('off')
        
        plt.show()
    
    # blobifier.malwareBlobs = blobifier.malwareBlobs[:300] # for testing! makes things much faster for clustering, analysis. FOR DEBUG ONLY
    
    # cluster blobs if desired
    # clusterer = SequentialClusterer()
    # clusterer.addCandidates(blobifier.malwareBlobs)
    # clusterer.addCandidates(blobifier.benignBlobs)
    # clusterer.cluster()
    

    # distance analysis if desired
    # distanceAnalysis(blobifier, mal_x, ben_x)

    
def distanceAnalysis(blobifier, mal_x, ben_x):
    # do the distance analysis in part 5 of paper
    # somewhat sketchy, but seems significant
    
    
    # To see if this works at all..
    # blobifier.malwareBlobs = new_mal_x[:100]
    # blobifier.benignBlobs = new_ben_x[:100]
    # mal_x = mal_x[:100]
    # ben_x = ben_x[:100]
    
    for n in range(0, len(blobifier.malwareBlobs)):
        # correlation
        # blobifier.malwareBlobs[n] = np.array(blobifier.malwareBlobs[n]).flatten()
        
        # tf dif
        blobifier.malwareBlobs[n] = np.array(blobifier.malwareBlobs[n]).flatten()
        blobifier.malwareBlobs[n] = " ".join([str(item) for item in blobifier.malwareBlobs[n]])
        
        # substrings
        # flat = np.array(blobifier.malwareBlobs[n]).flatten()
        # blobifier.malwareBlobs[n] = np.trim_zeros(flat).tolist()
        
        # for n grams manually
        # blobifier.malwareBlobs[n] = [str(int(item)) for item in blobifier.malwareBlobs[n]]
        
    for n in range(0, len(blobifier.benignBlobs)):
        # correlation
        # blobifier.benignBlobs[n] = np.array(blobifier.benignBlobs[n]).flatten()
        
        # tf dif
        blobifier.benignBlobs[n] = np.array(blobifier.benignBlobs[n]).flatten()
        blobifier.benignBlobs[n] = " ".join([str(item) for item in blobifier.benignBlobs[n]])
    
        # substrings
        # flat = np.array(blobifier.benignBlobs[n]).flatten()
        # blobifier.benignBlobs[n] = np.trim_zeros(flat).tolist()
        
        # for n grams manually
        # blobifier.benignBlobs[n] = [str(int(item)) for item in blobifier.benignBlobs[n]]
        
    print("beginning test")
    
    
    
    # transform regular files like above
    for n in range(0, len(mal_x)):
        # tf dif
        mal_x[n] = np.array(mal_x[n]).flatten()
        mal_x[n] = " ".join([str(item) for item in mal_x[n]])
        
    for n in range(0, len(ben_x)):
        ben_x[n] = np.array(ben_x[n]).flatten()
        ben_x[n] = " ".join([str(item) for item in ben_x[n]])
        
    lowerN = 1
    upperN = 1
        
    print("performing tfidf")
        
    tfidf = TfidfVectorizer(ngram_range=(lowerN,upperN)).fit_transform(blobifier.malwareBlobs)
    similarity_matrix = tfidf * tfidf.T
    indices = np.triu_indices(similarity_matrix.shape[0], k=1)
    similarities = similarity_matrix[indices].flatten()
    print("Blobbed malware to malware")
    print("mean", np.mean(similarities), "max", np.max(similarities), "min", np.min(similarities), "std", np.std(similarities))
    
    benignTFIDF = TfidfVectorizer(ngram_range=(lowerN,upperN)).fit(blobifier.malwareBlobs)
    benignTFIDF = benignTFIDF.transform(blobifier.benignBlobs)
    crossSimilarities = tfidf * benignTFIDF.T
    crossSimilarities = crossSimilarities.A.flatten()
    print("Blobbed malware to benign")
    print("mean", np.mean(crossSimilarities), "max", np.max(crossSimilarities), "min", np.min(crossSimilarities), "std", np.std(crossSimilarities))
    
    benign_similarity_matrix = benignTFIDF * benignTFIDF.T
    indices = np.triu_indices(benign_similarity_matrix.shape[0], k=1)
    benign_similarities = benign_similarity_matrix[indices].flatten()
    print("Blobbed benign to benign")
    print("mean", np.mean(benign_similarities), "max", np.max(benign_similarities), "min", np.min(benign_similarities), "std", np.std(benign_similarities))
    
    postMalMalSim = similarities
    postMalBenSim = crossSimilarities
    postBenBenSim = benign_similarities
    
    
    print()
    print()
    
    tfidf = TfidfVectorizer(ngram_range=(lowerN,upperN)).fit_transform(mal_x)
    similarity_matrix = tfidf * tfidf.T
    indices = np.triu_indices(similarity_matrix.shape[0], k=1)
    similarities = similarity_matrix[indices].flatten()
    print("Normal malware to malware")
    print("mean", np.mean(similarities), "max", np.max(similarities), "min", np.min(similarities), "std", np.std(similarities))
    
    benignTFIDF = TfidfVectorizer(ngram_range=(lowerN,upperN)).fit(mal_x)
    benignTFIDF = benignTFIDF.transform(ben_x)
    crossSimilarities = tfidf * benignTFIDF.T
    crossSimilarities = crossSimilarities.A.flatten()
    print("Normal malware to benign")
    print("mean", np.mean(crossSimilarities), "max", np.max(crossSimilarities), "min", np.min(crossSimilarities), "std", np.std(crossSimilarities))
    
    
    benign_similarity_matrix = benignTFIDF * benignTFIDF.T
    indices = np.triu_indices(benign_similarity_matrix.shape[0], k=1)
    benign_similarities = benign_similarity_matrix[indices].flatten()
    print("Normal benign to benign")
    print("mean", np.mean(benign_similarities), "max", np.max(benign_similarities), "min", np.min(benign_similarities), "std", np.std(benign_similarities))
    
    preMalMalSim = similarities
    preMalBenSim = crossSimilarities
    preBenBenSim = benign_similarities
    
    print()
    print()
    print("Statistics test")
    
    # from scipy.stats import ttest_ind
    # print("Malware to Malware", ttest_ind(postMalMalSim.T, preMalMalSim.T, equal_var=False))
    # print("Malware to Benign", ttest_ind(postMalBenSim.T, preMalBenSim.T, equal_var=False))
    # print("Benign to Benign", ttest_ind(postBenBenSim.T, preBenBenSim.T, equal_var=False))

    from scipy import stats
    print("Malware to Malware", stats.ttest_rel(postMalMalSim.T, preMalMalSim.T))
    print("Malware to Benign", stats.ttest_rel(postMalBenSim.T, preMalBenSim.T))
    print("Benign to Benign", stats.ttest_rel(postBenBenSim.T, preBenBenSim.T))
    print("Malware to Malware kruskal", stats.kruskal(postMalMalSim.T, preMalMalSim.T))
    print("Malware to Benign kruskal", stats.kruskal(postMalBenSim.T, preMalBenSim.T))
    print("Benign to Benign kruskal", stats.kruskal(postBenBenSim.T, preBenBenSim.T))


    # https://stackoverflow.com/questions/44862712/td-idf-find-cosine-similarity-between-new-document-and-dataset
    # https://stackoverflow.com/questions/6255835/cosine-similarity-and-tf-idf?rq=1
    # https://github.com/scipy/scipy/issues/7759
    # https://www.itl.nist.gov/div898/handbook/prc/section4/prc41.htm

    sns.distplot(preMalBenSim)
    plt.show()
    sns.distplot(postMalBenSim)
    plt.show()
    sns.distplot(preBenBenSim)
    plt.show()
    sns.distplot(postBenBenSim)
    plt.show()
    sns.distplot(postMalMalSim)
    plt.show()
    sns.distplot(preMalMalSim)
    plt.show()
    
    # '''
    
    
    
    
    

def blobToImage(blob):
    tokens = blob.split(" ")
    # for token in tokens:
        # print(token)
        # print(token[:-2])
    # comes as 237.0, cut off .-0
    tokens = [int(token[:-2]) for token in tokens]
    arr = np.array(tokens)
    
    arr = arr.reshape((int(arr.shape[0]/25), 25))
    
    return arr
    
def lookAtClusters():
    # clusterFile = r".\clusters_sequential.pklz"
    clusterFile = r".\clusters_sequential_BENIGN.pklz"
    
    readMe = gzip.open(clusterFile, "r")
    
    clusters = pickle.load(readMe)
    
    
    print(len(clusters))
    # print(clusters[0][0])
    
    for cluster in clusters:
        rn.shuffle(cluster)
        firstFew = cluster[:20]
        
        fig = plt.figure()
        for n, blob in enumerate(firstFew):
            ax1 = fig.add_subplot(4,5,n+1)
            ax1.imshow(blobToImage(blob), cmap='gray')
            plt.axis('off')
        
        plt.show()
        
        
        
    
    
    
if __name__ == "__main__":
    print("--------------------Performing Salience Test--------------------")
    salienceTest(model, "")
    
    
    # print("--------------------Looking At Clusters--------------------")
    # lookAtClusters()
    
    print("--------------------Testing Over--------------------")