from scapy.all import *
from loguru import logger


def split_pcap(file_path: str, chunk_size: int, save_base_path: str = None):
    packets = PcapReader(file_path)
    chunk = []
    counter = 1
    for packet in packets:
        # logger.info(packet.time)
        chunk.append(packet)
        if len(chunk) == chunk_size:
            wrpcap(f'{save_base_path}/chunk_{counter}.pcap', chunk)
            chunk = []
            logger.info(f'chunk_{counter}.pcap saved')
            counter += 1

    if chunk:
        wrpcap(f'{save_base_path}/chunk_{counter}.pcap', chunk)


def get_packet_time(pkt: Packet):
    return pkt.time


if __name__ == '__main__':
    from utils.files import create_dir
    create_dir('../_dataset/pcap/Friday-WorkingHours')
    split_pcap('../_dataset/pcap/Friday-WorkingHours.pcap', 10000, '../_dataset/pcap/Friday-WorkingHours')