split.py

import os
import subprocess
import math

# Paths and constants
PCAP_UTIL_PATH = "/mnt/g/UtilsCode/pcaputils/build/pcap-pckt-split"
INPUT_DIR = "/mnt/f/NetworkFlowDataset/CAIDA"
OUTPUT_DIR_2018 = "/mnt/g/UtilsCode/pcaputils/splitdataset/caida2018_split"
OUTPUT_DIR_2019 = "/mnt/g/UtilsCode/pcaputils/splitdataset/caida2019_split"
PCAP_FILES = {
    "caida2018.pcap": 22458659,  # Total packets from pcap-read-all
    "caida2019.pcap": 29922873
}
NUM_SPLITS = 10

def create_output_directory(output_dir):
    """Create output directory if it doesn't exist."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created directory: {output_dir}")

def split_pcap(pcap_file, total_packets, output_dir):
    """Split a pcap file into NUM_SPLITS parts with approximately equal packets."""
    packets_per_split = math.ceil(total_packets / NUM_SPLITS)
    print(f"Splitting {pcap_file}: {total_packets} packets into {NUM_SPLITS} parts, ~{packets_per_split} packets each")

    input_path = os.path.join(INPUT_DIR, pcap_file)

    for i in range(NUM_SPLITS):
        # Calculate 1-based start and end indices
        start_idx = i * packets_per_split + 1
        end_idx = min((i + 1) * packets_per_split, total_packets)
        if start_idx > total_packets:
            break

        # Generate output file name (e.g., caida2018_part_001.pcap)
        output_file = os.path.join(output_dir, f"{pcap_file.replace('.pcap', '')}_part_{i+1:03d}.pcap")

        # Construct the command for pcap-pckt-split
        cmd = [
            PCAP_UTIL_PATH,
            input_path,
            str(start_idx),
            str(end_idx),
            output_file
        ]
        print(f"Running command: {' '.join(cmd)}")
        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            print(f"Successfully created {output_file}")
        except subprocess.CalledProcessError as e:
            print(f"Error splitting {pcap_file} (part {i+1}): {e.stderr}")
            raise

def main():
    # Create output directories
    create_output_directory(OUTPUT_DIR_2018)
    create_output_directory(OUTPUT_DIR_2019)

    # Split each pcap file
    for pcap_file, total_packets in PCAP_FILES.items():
        output_dir = OUTPUT_DIR_2018 if "2018" in pcap_file else OUTPUT_DIR_2019
        split_pcap(pcap_file, total_packets, output_dir)

if __name__ == "__main__":
    main()