
# coding: utf-8

# In[48]:


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as I
import re
from scapy.all import *
import inspect
import pickle
import os
import numpy as np
import packet_preprocessing as pp
from packet_preprocessing import *
from train_dynamic import EmbedAttention, AttentionalBiRNN, HAN, tuple_batch
from torch.utils.data import DataLoader
# Default word tokens
PAD_token = 0  # Used for padding short sentences
SOS_token = 1  # Start-of-sentence token
EOS_token = 2  # End-of-sentence token
BD_token  = 3  # binary data token


# In[93]:


def readPackets(datafile):
    print("Reading packets...")
    packets = rdpcap(datafile+'/dump_sorted.pcap')
    print("Packets loaded")
    count = 0
    flag = False
    dnsrr_pack = []
    suspicious_pack = []
    
    for packet in packets:
        count += 1
        if packet.haslayer(DNSRR): 
            count = 0
            flag = True
            if isinstance(packet.an,DNSRR):
#                 print(packet.an.rrname)
                dnsrr_pack.append((packet,count))
        if flag == True:
            suspicious_pack.append(packet)
        if count == 200 and flag==True:
            break
    parsed_pack = [normalizeString(x.payload.command()) for x in suspicious_pack]
    return parsed_pack


# In[94]:


def preparedata(data_path):
    all_pack = []
    all_embedding = []
    print("Start preparing testing data ...")
    parsed_pack = readPackets(data_path)
    all_pack.append(parsed_pack)
    print("loaded: "+data_path)
    print("Counting fields...")
    print("Embedding...")
    for parsed_pack in all_pack:
        embedding = [indexesFromPacket(voc,packet) for packet in parsed_pack]
        all_embedding.append(embedding)
    return all_pack, all_embedding


# In[57]:


def testing_batch(review):
    list_rev = review


    sorted_r = sorted([(len(r),r_n,r) for r_n,r in enumerate(list_rev)],reverse=True) #index by desc rev_le
    lr,r_n,ordered_list_rev = zip(*sorted_r)
    lr = list(lr)
   
    max_sents = lr[0]


    #reordered
    review = [review[x] for x in r_n]

    stat =  sorted([(len(s),r_n,s_n,s) for r_n,r in enumerate(ordered_list_rev) for s_n,s in enumerate(r)],reverse=True)
    max_words = stat[0][0]


    ls = []
    batch_t = torch.zeros(len(stat),max_words).long()                         # (sents ordered by len)
    sent_order = torch.zeros(len(ordered_list_rev),max_sents).long().fill_(0) # (rev_n,sent_n)

    for i,s in enumerate(stat):
        sent_order[s[1],s[2]] = i+1 #i+1 because 0 is for empty.
        batch_t[i,0:len(s[3])] = torch.LongTensor(s[3])
        ls.append(s[0])

    return batch_t,sent_order,ls,lr,review


# In[96]:


if __name__ == '__main__':
    with open('f2n.pkl','rb') as f:
        fam2n = pickle.load(f)
    data_path = sys.argv[1]
#    data_path = '26'
    Dict = torch.load("test_save")
    voc = Dict["word_dic"] 
    Dict.popitem("word_dic")
    p, e = preparedata(data_path)
    dtl = []
    for i in range(len(e)):
        dtl.append(e[i])
    dataloader = DataLoader(dtl, collate_fn=testing_batch,pin_memory=True)
    device = torch.device("cpu")
    for iteration, (batch_t,sent_order,ls,lr,review) in enumerate(dataloader):
        data = (batch_t,sent_order)
        data = list(map(lambda x:x.to(device),data))


    model = HAN(ntoken=len(voc.index2field), emb_size=200,hid_size=100, num_class=5)
    model.load_state_dict(Dict)
    predict = model(data[0],data[1],ls,lr)
    i = predict[0].data.numpy().argmax()
    if i == 0:
        print("This is a benign app.")
    else:
        family = fam2n.n2family[i]
        print("Attention! This app is most likely a malware classified as ",family,"!")

