ID2T-toolkit-BotnetTraffic/code/ID2TLib/PaddingGenerator.py

from scapy.packet import Raw
import sys

import PayloadGenerator


def add_padding(packet, bytes_padding = 0):
    '''
    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
    :return: the initial packet, extended with the wanted amount of bytes of padding
    '''
    if(bytes_padding > 100):
        bytes_padding = 100
    payload = PayloadGenerator.generate_Payload(bytes_padding)
    packet = packet / Raw(load = payload)
    return packet

def equal_length(list_of_packets, length = 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 = sys.getsizeof(packet)
        if(packet_length > length):
            largest_packet = packet_length

    for packet in list_of_packets:
        bytes_padding = largest_packet - sys.getsizeof(packet)
        if(bytes_padding > 0):
            add_padding(packet, bytes_padding)

    return list_of_packets