"""
PCAP READER Module
~~~~~~~~~~~~~~~~~~
Author: Anish Gupta

Contains code that will convert Packet data in PCAP files to format acceptable to the ANN.
"""

from scapy.all import *
from datasketch.minhash import MinHash
from enum import Enum
import numpy as np
import json#,ipaddress
from scapy.layers import http

from hash import calculate_hash, entropy_domain_names


class SessionObjectOrder(Enum):
    """

    """
    PROTOCOL = 0
    IP1 = 1
    PORT1 = 2
    IP2 = 3
    PORT2 = 4
    HASH = 5


_MINHASH_INDEX = 4


class PcapReader:
    """
    PCAP Reader class.
    """
    def __init__(self,malware,num_pcap,filename=None):
        """
        :param filename: reads file dumped with dump_data (default: None)
        """
        self.num_pcap = num_pcap
        self.malware = malware
        self.session_collection = {}
        self.malware_feature_list = list()
        if filename is not None:
            dump_file = open(filename, 'r')
            data = dump_file.read()
            data = json.loads(data)
            for key in data:
                hash_val = np.asarray(data[key][5], dtype=np.uint64)
                data[key][_MINHASH_INDEX] = MinHash()
                data[key][_MINHASH_INDEX] = hash_val

    def read_pcap(self, filename):
        """
        Reads PCAP files using scapy's rdpcap, divides data into sessions and creates a hash for each session, a hash
        can be updated if sessions exist in subsequent files.

        :param filename: PCAP file path
        """
        packets = rdpcap(filename)[IP]
        sessions = packets.sessions()
        for key in sessions:
            # try:
            if key not in self.session_collection:
                parts = key.split()
                protocol = parts[0]
                ip1_parts = parts[1].split(':')
                ip1 = ip1_parts[0]
                port1 = int(ip1_parts[1]) if len(ip1_parts) > 1 else 0
                ip2_parts = parts[3].split(':')
                ip2 = ip2_parts[0]
                port2 = int(ip2_parts[1]) if len(ip1_parts) > 1 else 0
                entropy = entropy_domain_names(sessions[key])
                self.session_collection[key] = [
                    hash(protocol),
                    entropy,
                    # int(ipaddress.ip_address(ip1)),
                    port1,
                    # int(ipaddress.ip_address(ip2)),
                    port2,
                    MinHash()
                ]
            self.session_collection[key] = calculate_hash(self.session_collection[key], sessions[key])
            # except Exception as e:
            #     print(str(e))
            #     pass

    def read_malware(self):
        for i in range(1,self.num_pcap):
            filename = "./data/" + self.malware + "0" + str(i) + ".pcap"
            print(filename)
            self.read_pcap(filename)
            self.malware_feature_list.append(self.get_data())
        if len(self.malware_feature_list) > 0:
            self.malware_feature_list = np.concatenate(self.malware_feature_list)
        return self.malware_feature_list

    def get_data(self):
        """
        Get all the data as a list.

        Format
        ~~~~~~

        PROTOCOL = 0
        IP1 = 1
        PORT1 = 2
        IP2 = 3
        PORT2 = 4
        HASH = 5

        :return: List of all data
        """
        response = []
        for key in self.session_collection:
            sess = self.session_collection[key]
            session_object = sess[0:_MINHASH_INDEX]
            hash_array = sess[_MINHASH_INDEX].hashvalues
            for num in hash_array:
            	session_object.append(num)
            response.append(session_object)
        return response

    def dump_data(self, filename):
        """
        Dumps data into a JSON file.

        :param filename: Filename of the dump file
        """
        response = {}
        for key in self.session_collection:
            sess = self.session_collection[key]
            session_object = sess[0:_MINHASH_INDEX]
            session_object.append(sess[_MINHASH_INDEX].hashvalues.tolist())
            response[key] = session_object
        dump_file = open(filename, 'w')
        dump_file.write(json.dumps(response))
        dump_file.close()
