
import pandas as pd
import numpy as np
from sklearn import preprocessing
from keras.utils.np_utils import to_categorical
from keras.layers import Masking, Embedding
from keras.layers import LSTM
le = preprocessing.LabelEncoder()
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
sequenceLength = -1
dirname = '/home/malware/trace1_test/'

def loadTest():
    global sequenceLength
    # label = pd.read_csv(dirname+'safetype.txt',sep=';')
    train = pd.read_csv(dirname+'sandbox_behaviorlist.txt',sep=';')
    shaList = train['sha1']
    result = []
    for behaviour in train['behavior_list']:
        if type(behaviour) == float:
            []
        else:
            bList = behaviour.split(',')
            sequenceLength = max(sequenceLength,len(bList))
            tempList = [int(x) for x in bList]
        result.append(tempList)
    for belist in result:
        nowlen = len(belist)
        if nowlen < sequenceLength:
            belist += [-1]*(sequenceLength-nowlen)
    X = np.array(result).reshape(len(result),sequenceLength,1)
    return X,shaList

from keras.models import load_model
X,shaList = loadTest()
model = load_model('test.h5')
y_pred = model.predict(X)
outf = open('data/lstm_output.txt','w+')
outf.write('sha1,0,1,2,3,4,5\n')
for i in range(len(shaList)):
    result = y_pred[i]
    outf.write('%s,%f,%f,%f,%f,%f,%f\n'%(shaList[i],result[0],result[1],result[2],result[3],result[4],result[5]))
