
# coding: utf-8

# In[538]:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals


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
# 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[507]:


class f2n():
    def __init__(self):
        self.family2n = {}
        self.file2n = {}
        self.n2family = {}
        self.num_fam = 0
        self.load_family()
    def add_family(self,fam):
        if fam not in self.family2n:
            self.family2n[fam] = self.num_fam
            self.n2family[self.num_fam] = fam
            self.num_fam+=1
    def load_family(self):
        for root, dirs, files in os.walk("./malware/", topdown=False):
            for name in files:
                family = root.split('/')[2]
                if name.endswith('icloud'):
                    name = name.split('.')[1]+'.'+name.split('.')[2]
                if len(family):
                    self.add_family(family)
                    if name not in self.file2n:
                        self.file2n[name] = self.family2n[family]


# In[519]:


def get_packet_family(directory,fam2n):
    with open(directory+'/analysis.log') as f:
        file_name = f.readlines()[2].split(' ')[-1].split('/')[-1].rstrip()
        family_num = fam2n.file2n[file_name]
        return file_name,family_num


# In[513]:


class Voc:
    def __init__(self):
        self.field2index = {}
        self.field2count = {"BD":0}
        self.index2field = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS",BD_token: "BD"}
        self.num_fields = 4  # Count SOS, EOS, PAD, BD
        self.trimmed = False

    def addPacket(self, packet):
        for field in packet.split(' '):
            self.addfield(field)
    def search_field_index(self, field):
        if len(field) == 3 and field[0]=='x':
            return 3
        else:
            try:
                a = self.field2index[field]
                return a
            except:
                return 3


    def addfield(self, field):
        if field not in self.field2index:
            if len(field) == 3 and field[0]=='x':
                self.field2count["BD"]+=1
            else:
                self.field2index[field] = self.num_fields
                self.field2count[field] = 1
                self.index2field[self.num_fields] = field
                self.num_fields += 1
        else:
            self.field2count[field] += 1
    def trim(self, min_count):
        if self.trimmed:
            return
        self.trimmed = True
        keep_fields = []
        for k, v in self.field2count.items():
            if v >= min_count:
                keep_fields.append(k)

        print('keep_fields {} / {} = {:.4f}'.format(
            len(keep_fields), len(self.field2index), len(keep_fields) / len(self.field2index)
        ))
        # Reinitialize dictionaries
        self.field2index = {}
        self.field2count = {}
        self.index2field = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS",BD_token: "BD"}
        self.num_fields = 4 # Count default tokens
        for field in keep_fields:
            self.addfield(field)

    def load_data(self, packet):
        data = load_preprocess()
        for packet in data[0]:
            normalizeString(packet.command())
# Lowercase and remove non-letter characters
def normalizeString(s):
    s = re.sub(r"[^a-zA-Z0-9.!?]+", r" ", s)
    return s


# Takes string pacet.command(), returns packet of field indexes
def indexesFromPacket(voc, packet):
    return [voc.search_field_index(field) for field in packet.strip().split(' ')] + [EOS_token]


# In[539]:


def readPackets(datafile,fam2n=None):
    print("Reading packets...")
    print(datafile)
    packets = rdpcap(datafile+'/dump_sorted.pcap')
    if not fam2n:
        fam2n = f2n()
    file_name, fam_n = get_packet_family(datafile,fam2n)
    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 == 50 and flag==True:
            break
    parsed_pack = [normalizeString(x.payload.command()) for x in suspicious_pack]
    return parsed_pack, fam_n, file_name
# voc, parsed_pack, fam_n = readPackets(datafile)
# print(parsed_pack[0])
# print(fam_n)


# In[535]:


def preparedata(files,fam2n=None):
    voc = Voc()
    all_pack = []
    all_fam = []
    all_embedding = []
    count = 0
    print("Start preparing training data ...")
    for file in files:
        #if count == 5:
        #    break
        #count += 1
        try:
            parsed_pack, fam_n, file_name = readPackets(file,fam2n)
            all_pack.append(parsed_pack)
            all_fam.append(fam_n)
            print("loaded: "+file_name)
        except:
            pass
    print("Counting fields...")
    for parsed_pack in all_pack:
        for packet in parsed_pack:
            voc.addPacket(packet)
    voc.trim(10)
    print("Counted fields:", voc.num_fields)
    print("Embedding...")
    for parsed_pack in all_pack:
        embedding = [indexesFromPacket(voc,packet) for packet in parsed_pack]
        all_embedding.append(embedding)
    return voc, all_pack, all_fam, all_embedding
if __name__ == '__main__':
    data_path = sys.argv[1]
    datafile = next(os.walk(data_path))[1]
    with open('f2n.pkl','rb') as f:
         fam2n = pickle.load(f)
    voc, all_pack, all_fam,all_embedding = preparedata(datafile,fam2n)
    with open('embedding.pkl', 'wb') as f:
        pickle.dump(all_embedding,f)
        f.close()
    with open('fam.pkl','wb') as f:
        pickle.dump(all_fam,f)
        f.close()
    with open('voc.pkl','wb') as f:
        pickle.dump(voc,f)
        f.close()
    

