from scapy.packet import Raw

import ID2TLib.PayloadGenerator as PayloadGenerator
import numpy.random as random


def add_padding(packet, bytes_padding = 0, user_padding=True, rnd = False):
    '''
    Adds padding to a packet with the given amount of bytes, but a maximum of 100 bytes.
    :param packet: the packet that will be extended with the additional payload
    :param bytes_padding: the amount of bytes that will be appended to the packet
    :param user_padding: true, if the function add_padding is called from another class and the user
    sets the padding manually
    :param rnd: adds a random padding betwing 0 and bytes_adding, if true
    :return: the initial packet, extended with the wanted amount of bytes of padding
    '''

    if(user_padding and bytes_padding > 100):
        bytes_padding = 100

    if (rnd is True):
        r = int(round(bytes_padding / 4))                  #sets bytes_padding to any number between 0 and bytes_padding
        bytes_padding = random.random_integers(0, r) * 4   #, that's dividable by 4
    payload = PayloadGenerator.generate_payload(bytes_padding)
    packet[Raw].load += Raw(load=payload).load
    return packet

def equal_length(list_of_packets, length = 0, padding = 0):
    '''
    Equals the length of a given set of packets on the given length. If the given length is smaller than the largest
    packet, all the other packets are extended to the largest packet's length.
    :param list_of_packets: The given set of packet.
    :param length: The length each packet should have.
    :return: The set of extended packets.
    '''

    largest_packet = length
    for packet in list_of_packets:
        packet_length = len(packet)
        if(packet_length > largest_packet):
            largest_packet = packet_length

    for packet in list_of_packets:
        bytes_padding = largest_packet - len(packet)
        if(bytes_padding > 0):
            add_padding(packet, bytes_padding, False, False)
            add_padding(packet, padding, False, True)

    return list_of_packets