from scapy.all import *
from scapy.layers.inet import IP, UDP, TCP
from collections import defaultdict

from scapy.plist import PacketList


def _load_pcap(file_name: str) -> PacketList:
    pkts = rdpcap(file_name)
    return pkts

def  _filename_gen(t:tuple):
    proto = "UNKNOWN"
    if t[0] == 6:
        proto = "TCP"
    if t[0] == 17:
        proto = "UDP"
    return f"{proto}_{t[1]}_{t[3]}_{t[2]}_{t[4]}"

def process(packets: PacketList):
    five_tuple_classified = defaultdict(list)
    for pkt in packets:
        ip_layer = pkt[IP]
        if TCP in pkt:
            transmission_layer = pkt[TCP]
        elif UDP in pkt:
            transmission_layer = pkt[UDP]
        else:
            continue
        key = (ip_layer.proto, ip_layer.src, ip_layer.dst, transmission_layer.sport, transmission_layer.dport)
        five_tuple_classified[key].append(pkt)

    for key, value in five_tuple_classified.items():
        print(key, value)
        wrpcap(f"split_output/{_filename_gen(key)}.pcap", value)


# path = "./facebook_audio1a.pcap"
path = "./vpn_aim_chat1a.pcap"
pkts = _load_pcap(path)
process(pkts)