from datetime import datetime
import ipaddress

################################################################################
#                              Single Flow object                              #
################################################################################

class Flow(object):
    """Flow object extracted from pcap file that can be used for fingerprinting

        Attributes
        ----------
        src : string
            Source IP

        sport : int
            Source port

        dst : string
            Destination IP

        dport : int
            Destination port

        source : tuple
            (Source IP, source port) tuple

        destination : tuple
            (Destination IP, destination port) tuple

        time_start : int
            Timestamp of first packet in flow

        time_end : int
            Timestamp of last packet in flow

        certificate : Object
            Certificate of flow, if any

        ip_lengths : list
            List of packet length for each ip packet in flow
        payload_lengths : list
            List of payload sequence for each tcp/udp fragment with non-zero payload in flow.
        ip_timestamps : list
            List of timestamps corresponding to each ip packet in flow, it may contain packets without any tcp/udp payload.

        payload_timestamps: list
            List of timestamps corresponding to each tcp/udp fragment with non-zero payload in flow.
    """

    def __init__(self,main='payload'):
        """
            param
            -----------
            main: str
              'payload' means the main lengths sequence and timestampes sequence refer to packets with non-zero payload, the sequences will fitler out zero payload packets.
              'ip'   means the main lengths sequence and timestamps sequence refer to any packets, it will not filter any packets.
        """
        self.main = main
        """Initialise an empty Flow."""
        # Initialise flow endpoints
        self.src   = None
        self.sport = None
        self.dst   = None
        self.dport = None

        # Initialise extension
        self.extension = dict()
        # Initialise packet lengths
        self.ip_lengths   = list()
        self.payload_lengths = list()
        # Initialise packet timestamps
        self.ip_timestamps = list()
        self.payload_timestamps = list()   #non-zero payload packet's timestamp sequence.

        # Refer the main property
        self.lengths = self.payload_lengths if main=='payload' else self.ip_lengths
        self.timestamps = self.payload_timestamps if main=='payload' else self.ip
    ########################################################################
    #                        Add new packet to flow                        #
    ########################################################################

    def add(self, packet,extension):
        """Add a new packet to the flow.

            Parameters
            ----------
            packet : np.array of shape=(n_features,)
                Packet from Reader.

            Returns
            -------
            self : self
                Returns self
            """
        #print(packet)
        try:
            # Extract IPs from packet
            ip_a, ip_b = packet[5], packet[6]
        except BaseException as exp:
            raise ValueError('This is not ip packet.')
        try:
            # Extract ports from packet
            port_a, port_b = int(packet[7]), int(packet[8])
        except BaseException as exp:
            raise ValueError('This ip packet is not sample of tcp or udp or gre.')
        # Perform packet check
        if self.src is not None:
            if {self.src, self.dst} != {ip_a, ip_b} and {self.sport, self.dport} != {port_a, port_b}:
                raise ValueError("Packet {} incompatible with flow {}" .format(packet, self))
        # Set endpoints where smallest dport is destination
        elif port_a > port_b:
            self.src  , self.dst   = ip_a  , ip_b
            self.sport, self.dport = port_a, port_b
        else:
            self.src  , self.dst   = ip_b  , ip_a
            self.sport, self.dport = port_b, port_a

        # Add extension if any
        for i in range(len(packet[-1])):
            #if packet[-1][i] != "":
                if extension[i] not in self.extension:
                    self.extension.setdefault(extension[i],[])
                self.extension[extension[i]].append(packet[-1][i])


        # Set timestamps and lengths
        #print(packet)
        self.ip_timestamps.append(float(packet[3]))
        self.ip_lengths   .append( int(packet[4]) if packet[5] == self.src else
                               -int(packet[4]))
        if int(packet[9]) != 0:
            self.payload_lengths.append( int(packet[9]) if packet[5] == self.src else
                               -int(packet[9]))
            self.payload_timestamps.append(float(packet[3]))

        # Return self
        return self

    ########################################################################
    #                  Source/Destination/Time attributes                  #
    ########################################################################

    @property
    def source(self):
        """(source IP, source port)-tuple of Flow"""
        return (self.src, self.sport)

    @property
    def destination(self):
        """(destination IP, destination port)-tuple of Flow"""
        return (self.dst, self.dport)

    @property
    def time_start(self):
        """Returns start time of Flow"""
        return min(self.timestamps)

    @property
    def time_end(self):
        """Returns end time of Flow"""
        return max(self.timestamps)

    ########################################################################
    #                           Class overrides                            #
    ########################################################################

    def __len__(self):
        """Return length of Flow in packets."""
        return len(self.lengths)

    def __str__(self):
        """Return string representation of flow."""
        if self.main=='ip':
            return "[Time {} to {}] {:>15}:{:<5} <-> {:>15}:{:<5} [IP Packet Size Length {}] [SNI: {}]".format(
                datetime.fromtimestamp(min(self.timestamps)).strftime("%H:%M:%S.%f"),
                datetime.fromtimestamp(max(self.timestamps)).strftime("%H:%M:%S.%f"),
                self.src, self.sport, self.dst, self.dport,
                len(self),self.extension)
        else:
            return "[Time {} to {}] {:>15}:{:<5} <-> {:>15}:{:<5} [Payload Packet Size Length {}] [SNI: {}]".format(
                datetime.fromtimestamp(min(self.timestamps)).strftime("%H:%M:%S.%f"),
                datetime.fromtimestamp(max(self.timestamps)).strftime("%H:%M:%S.%f"),
                self.src, self.sport, self.dst, self.dport,
                len(self),self.extension)
    def __gt__(self, other):
        """Greater than object override"""
        return min(self.timestamps) >  min(other.timestamps)

    def __ge__(self, other):
        """Greater equals object override"""
        return min(self.timestamps) >= min(other.timestamps)

    def __lt__(self, other):
        """Less than object override"""
        return min(self.timestamps) <  min(other.timestamps)

    def __le__(self, other):
        """Less equals object override"""
        return min(self.timestamps) <= min(other.timestamps)