generate_test_pcap.py

#!/usr/bin/env python3
"""
Generate a large test PCAP file with various protocols for DPI testing.
Generates 1500+ flows with balanced classes for ML training.
"""

import struct
import random

class PCAPWriter:
    def __init__(self, filename):
        self.file = open(filename, 'wb')
        self.write_global_header()
        self.timestamp = 1700000000
        
    def write_global_header(self):
        header = struct.pack('<IHHIIII', 0xa1b2c3d4, 2, 4, 0, 0, 65535, 1)
        self.file.write(header)
        
    def write_packet(self, data):
        ts_sec  = self.timestamp
        ts_usec = random.randint(0, 999999)
        self.timestamp += random.randint(1, 3)
        pkt_header = struct.pack('<IIII', ts_sec, ts_usec, len(data), len(data))
        self.file.write(pkt_header)
        self.file.write(data)
        
    def close(self):
        self.file.close()


def create_ethernet_header(src_mac, dst_mac, ethertype=0x0800):
    return bytes.fromhex(dst_mac.replace(':', '')) + \
           bytes.fromhex(src_mac.replace(':', '')) + \
           struct.pack('>H', ethertype)


def create_ip_header(src_ip, dst_ip, protocol, payload_len):
    version_ihl = 0x45
    tos         = 0
    total_len   = 20 + payload_len
    ident       = random.randint(1, 65535)
    flags_frag  = 0x4000
    ttl         = 64
    checksum    = 0
    header = struct.pack('>BBHHHBBH',
                         version_ihl, tos, total_len,
                         ident, flags_frag, ttl, protocol, checksum)
    header += bytes([int(x) for x in src_ip.split('.')])
    header += bytes([int(x) for x in dst_ip.split('.')])
    return header


def create_tcp_header(src_port, dst_port, seq, ack, flags, payload_len=0):
    data_offset = 5 << 4
    window      = 65535
    checksum    = 0
    urgent      = 0
    return struct.pack('>HHIIBBHHH',
                       src_port, dst_port, seq, ack,
                       data_offset, flags, window, checksum, urgent)


def create_udp_header(src_port, dst_port, payload_len):
    length   = 8 + payload_len
    checksum = 0
    return struct.pack('>HHHH', src_port, dst_port, length, checksum)


def create_tls_client_hello(sni):
    sni_bytes  = sni.encode('ascii')
    sni_entry  = struct.pack('>BH', 0, len(sni_bytes)) + sni_bytes
    sni_list   = struct.pack('>H', len(sni_entry)) + sni_entry
    sni_ext    = struct.pack('>HH', 0x0000, len(sni_list)) + sni_list
    sup_ver    = struct.pack('>HHB', 0x002b, 3, 2) + struct.pack('>H', 0x0304)
    extensions = sni_ext + sup_ver
    ext_data   = struct.pack('>H', len(extensions)) + extensions
    client_ver = struct.pack('>H', 0x0303)
    rand_bytes = bytes([random.randint(0, 255) for _ in range(32)])
    session_id = struct.pack('B', 0)
    ciphers    = struct.pack('>H', 4) + struct.pack('>HH', 0x1301, 0x1302)
    compress   = struct.pack('BB', 1, 0)
    ch_body    = client_ver + rand_bytes + session_id + ciphers + compress + ext_data
    handshake  = struct.pack('B', 0x01) + struct.pack('>I', len(ch_body))[1:] + ch_body
    record     = struct.pack('B', 0x16) + struct.pack('>H', 0x0301) + \
                 struct.pack('>H', len(handshake)) + handshake
    return record


def create_tcp_fin(src_port, dst_port, seq, ack):
    data_offset = 5 << 4
    flags       = 0x11  # FIN + ACK
    window      = 65535
    return struct.pack('>HHIIBBHHH',
                       src_port, dst_port, seq, ack,
                       data_offset, flags, window, 0, 0)


def create_http_request(host, path='/'):
    return f"GET {path} HTTP/1.1\r\nHost: {host}\r\nUser-Agent: DPI-Test/1.0\r\nAccept: */*\r\n\r\n".encode()


def create_dns_query(domain):
    txid     = struct.pack('>H', random.randint(1, 65535))
    flags    = struct.pack('>H', 0x0100)
    counts   = struct.pack('>HHHH', 1, 0, 0, 0)
    question = b''
    for label in domain.split('.'):
        question += struct.pack('B', len(label)) + label.encode()
    question += struct.pack('B', 0)
    question += struct.pack('>HH', 1, 1)
    return txid + flags + counts + question


def write_tls_flow(writer, user_ip, dst_ip, sni, dst_port,
                   user_mac, gateway_mac, seq_base,
                   extra_data_packets=2, send_fin=True):
    src_port = random.randint(49152, 65535)

    # SYN
    eth = create_ethernet_header(user_mac, gateway_mac)
    tcp = create_tcp_header(src_port, dst_port, seq_base, 0, 0x02)
    ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp))
    writer.write_packet(eth + ip + tcp)

    # SYN-ACK
    eth = create_ethernet_header(gateway_mac, user_mac)
    tcp = create_tcp_header(dst_port, src_port, seq_base+1000, seq_base+1, 0x12)
    ip  = create_ip_header(dst_ip, user_ip, 6, len(tcp))
    writer.write_packet(eth + ip + tcp)

    # ACK
    eth = create_ethernet_header(user_mac, gateway_mac)
    tcp = create_tcp_header(src_port, dst_port, seq_base+1, seq_base+1001, 0x10)
    ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp))
    writer.write_packet(eth + ip + tcp)

    # TLS Client Hello with SNI
    tls = create_tls_client_hello(sni)
    tcp = create_tcp_header(src_port, dst_port, seq_base+1, seq_base+1001, 0x18)
    ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp) + len(tls))
    writer.write_packet(eth + ip + tcp + tls)

    # Data exchange packets
    seq = seq_base + 100
    for _ in range(extra_data_packets):
        payload = bytes([random.randint(0, 255) for _ in range(random.randint(100, 1400))])
        tcp = create_tcp_header(src_port, dst_port, seq, seq_base+1001, 0x18)
        ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp) + len(payload))
        writer.write_packet(eth + ip + tcp + payload)
        seq += len(payload)

    # FIN — closes flow and triggers CSV export
    if send_fin:
        tcp = create_tcp_fin(src_port, dst_port, seq, seq_base+1001)
        ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp))
        writer.write_packet(eth + ip + tcp)

    return seq_base + 10000


def write_http_flow(writer, user_ip, dst_ip, host, dst_port,
                    user_mac, gateway_mac, seq_base, send_fin=True):
    src_port = random.randint(49152, 65535)
    eth      = create_ethernet_header(user_mac, gateway_mac)

    tcp = create_tcp_header(src_port, dst_port, seq_base, 0, 0x02)
    ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp))
    writer.write_packet(eth + ip + tcp)

    http = create_http_request(host)
    tcp  = create_tcp_header(src_port, dst_port, seq_base+1, 1, 0x18)
    ip   = create_ip_header(user_ip, dst_ip, 6, len(tcp) + len(http))
    writer.write_packet(eth + ip + tcp + http)

    if send_fin:
        tcp = create_tcp_fin(src_port, dst_port, seq_base+200, 1)
        ip  = create_ip_header(user_ip, dst_ip, 6, len(tcp))
        writer.write_packet(eth + ip + tcp)

    return seq_base + 10000


def write_dns_flow(writer, user_ip, dns_ip, domain, user_mac, gateway_mac):
    src_port = random.randint(49152, 65535)
    dns_data = create_dns_query(domain)
    eth      = create_ethernet_header(user_mac, gateway_mac)
    udp      = create_udp_header(src_port, 53, len(dns_data))
    ip       = create_ip_header(user_ip, dns_ip, 17, len(udp) + len(dns_data))
    writer.write_packet(eth + ip + udp + dns_data)


def main():
    writer = PCAPWriter('test_dpi.pcap')

    user_mac    = '00:11:22:33:44:55'
    gateway_mac = 'aa:bb:cc:dd:ee:ff'
    dns_server  = '8.8.8.8'
    seq         = 1000

    # ----------------------------------------------------------------
    # TLS flows
    # Format: (dst_ip, sni, dst_port, extra_data_pkts, repeat)
    # repeat = how many times this flow is generated
    # Minimum 50 rows per app for good ML accuracy
    # ----------------------------------------------------------------
    TLS_FLOWS = [
        # ── YouTube (target: 80+ rows) ──────────────────────────────
        ('142.250.185.110', 'www.youtube.com',          443, 8,  25),
        ('142.250.185.120', 'youtu.be',                 443, 6,  20),
        ('172.217.14.78',   'googlevideo.com',          443, 10, 20),
        ('142.250.185.130', 'ytimg.com',                443, 4,  15),

        # ── Netflix (target: 80+ rows) ───────────────────────────────
        ('23.52.167.61',    'www.netflix.com',          443, 12, 25),
        ('23.52.167.70',    'nflxvideo.net',            443, 15, 25),
        ('54.155.178.5',    'nflximg.net',              443, 8,  25),

        # ── Facebook (target: 60+ rows) ──────────────────────────────
        ('157.240.1.35',    'www.facebook.com',         443, 4,  25),
        ('157.240.1.180',   'fbcdn.net',                443, 3,  25),

        # ── Instagram (target: 60+ rows) ─────────────────────────────
        ('157.240.1.174',   'www.instagram.com',        443, 5,  30),
        ('157.240.1.190',   'cdninstagram.com',         443, 4,  25),

        # ── WhatsApp (target: 60+ rows) ──────────────────────────────
        ('157.240.1.200',   'web.whatsapp.com',         443, 6,  30),
        ('157.240.1.210',   'whatsapp.net',             443, 5,  25),

        # ── Google (target: 80+ rows) ────────────────────────────────
        ('142.250.185.206', 'www.google.com',           443, 3,  20),
        ('142.250.185.100', 'mail.google.com',          443, 4,  15),
        ('142.250.185.150', 'drive.google.com',         443, 5,  15),
        ('142.250.185.160', 'docs.google.com',          443, 3,  15),
        ('142.250.185.170', 'googleapis.com',           443, 2,  15),

        # ── Twitter/X (target: 60+ rows) ─────────────────────────────
        ('104.244.42.65',   'twitter.com',              443, 3,  20),
        ('104.244.42.66',   'api.twitter.com',          443, 2,  20),
        ('104.244.42.67',   't.co',                     443, 2,  15),

        # ── Amazon (target: 60+ rows) ────────────────────────────────
        ('52.94.236.248',   'www.amazon.com',           443, 4,  25),
        ('52.94.236.100',   's3.amazonaws.com',         443, 6,  20),
        ('52.94.236.150',   'api.amazon.com',           443, 3,  15),

        # ── Microsoft (target: 60+ rows) ─────────────────────────────
        ('13.107.42.14',    'www.microsoft.com',        443, 3,  15),
        ('13.107.42.20',    'outlook.live.com',         443, 4,  15),
        ('13.107.42.30',    'teams.microsoft.com',      443, 5,  15),
        ('13.107.42.40',    'onedrive.live.com',        443, 6,  15),

        # ── Apple (target: 60+ rows) ─────────────────────────────────
        ('17.253.144.10',   'www.apple.com',            443, 3,  20),
        ('17.253.144.20',   'icloud.com',               443, 4,  20),
        ('17.253.144.30',   'api.apple.com',            443, 2,  20),

        # ── Spotify (target: 60+ rows) ───────────────────────────────
        ('35.186.224.47',   'open.spotify.com',         443, 6,  30),
        ('35.186.224.50',   'api.spotify.com',          443, 3,  25),

        # ── Discord (target: 60+ rows) ───────────────────────────────
        ('104.16.85.20',    'discord.com',              443, 4,  30),
        ('104.16.85.21',    'cdn.discordapp.com',       443, 5,  25),

        # ── Zoom (target: 60+ rows) ──────────────────────────────────
        ('35.186.224.25',   'zoom.us',                  443, 8,  30),
        ('35.186.224.26',   'api.zoom.us',              443, 3,  25),

        # ── Telegram (target: 60+ rows) ──────────────────────────────
        ('35.186.227.140',  'web.telegram.org',         443, 5,  30),
        ('149.154.167.50',  'api.telegram.org',         443, 4,  25),

        # ── TikTok (target: 60+ rows) ────────────────────────────────
        ('99.86.0.100',     'www.tiktok.com',           443, 7,  25),
        ('99.86.0.110',     'api.tiktok.com',           443, 4,  25),

        # ── GitHub (target: 60+ rows) ────────────────────────────────
        ('140.82.114.4',    'github.com',               443, 4,  20),
        ('140.82.114.5',    'api.github.com',           443, 3,  20),
        ('140.82.114.6',    'raw.githubusercontent.com',443, 5,  20),

        # ── Cloudflare (target: 60+ rows) ────────────────────────────
        ('192.0.78.24',     'www.cloudflare.com',       443, 3,  30),
        ('104.16.0.1',      'cdn.cloudflare.com',       443, 2,  25),
        ('104.16.1.1',      'dash.cloudflare.com',      443, 3,  10),
    ]

    HTTP_FLOWS = [
        ('93.184.216.34',   'example.com',              80,  20),
        ('185.199.108.153', 'httpbin.org',              80,  15),
        ('104.21.0.1',      'neverssl.com',             80,  15),
        ('178.62.0.100',    'old-site.com',             80,  10),
    ]

    DNS_DOMAINS = [
        ('www.google.com',          30),
        ('www.youtube.com',         25),
        ('www.facebook.com',        20),
        ('api.twitter.com',         20),
        ('www.netflix.com',         20),
        ('open.spotify.com',        20),
        ('discord.com',             20),
        ('zoom.us',                 20),
        ('web.telegram.org',        20),
        ('www.tiktok.com',          20),
        ('github.com',              20),
        ('www.amazon.com',          20),
        ('www.microsoft.com',       20),
        ('icloud.com',              15),
        ('www.apple.com',           15),
        ('www.cloudflare.com',      15),
        ('web.whatsapp.com',        15),
        ('www.instagram.com',       15),
    ]

    # Multiple user IPs — simulates real network with many devices
    user_ips = [
        '192.168.1.100',
        '192.168.1.101',
        '192.168.1.102',
        '192.168.1.103',
        '192.168.1.104',
        '192.168.1.105',
        '192.168.1.106',
    ]

    # ── Generate TLS flows ──────────────────────────────────────────
    tls_count = 0
    for dst_ip, sni, dst_port, extra_pkts, repeat in TLS_FLOWS:
        for _ in range(repeat):
            user_ip = random.choice(user_ips)
            pkts    = max(1, extra_pkts + random.randint(-1, 2))
            seq     = write_tls_flow(writer, user_ip, dst_ip, sni, dst_port,
                                     user_mac, gateway_mac, seq,
                                     extra_data_packets=pkts,
                                     send_fin=True)
            tls_count += 1

    # ── Generate HTTP flows ─────────────────────────────────────────
    http_count = 0
    for dst_ip, host, dst_port, repeat in HTTP_FLOWS:
        for _ in range(repeat):
            user_ip = random.choice(user_ips)
            seq     = write_http_flow(writer, user_ip, dst_ip, host, dst_port,
                                      user_mac, gateway_mac, seq,
                                      send_fin=True)
            http_count += 1

    # ── Generate DNS queries ────────────────────────────────────────
    dns_count = 0
    for domain, repeat in DNS_DOMAINS:
        for _ in range(repeat):
            user_ip = random.choice(user_ips)
            write_dns_flow(writer, user_ip, dns_server, domain,
                           user_mac, gateway_mac)
            dns_count += 1

    # ── Blocked IP traffic ──────────────────────────────────────────
    # Run dpi_engine with: --block-ip 192.168.1.50
    blocked_ip  = '192.168.1.50'
    blocked_mac = '00:11:22:33:44:56'
    blocked_targets = [
        ('172.217.0.100', 'www.youtube.com',   443, 3),
        ('157.240.1.35',  'www.facebook.com',  443, 3),
        ('157.240.1.174', 'www.instagram.com', 443, 3),
        ('99.86.0.100',   'www.tiktok.com',    443, 3),
        ('23.52.167.61',  'www.netflix.com',   443, 3),
    ]
    blocked_count = 0
    for dst_ip, sni, dst_port, extra_pkts in blocked_targets:
        for _ in range(15):
            seq = write_tls_flow(writer, blocked_ip, dst_ip, sni, dst_port,
                                 blocked_mac, gateway_mac, seq,
                                 extra_data_packets=extra_pkts,
                                 send_fin=True)
            blocked_count += 1

    writer.close()

    total = tls_count + http_count + dns_count + blocked_count
    print(f"Created test_dpi.pcap")
    print(f"  TLS flows   : {tls_count}")
    print(f"  HTTP flows  : {http_count}")
    print(f"  DNS queries : {dns_count}")
    print(f"  Blocked IP  : {blocked_count}")
    print(f"  ─────────────────────────────")
    print(f"  Total flows : {total}  (~{total} rows expected in flows.csv)")
    print()
    print("Next steps:")
    print("  1. Run DPI engine to generate flows.csv:")
    print("       Windows : .\\build\\Release\\dpi_mt.exe test_dpi.pcap output.pcap")
    print("       Linux   : ./build/dpi_mt test_dpi.pcap output.pcap")
    print()
    print("  2. For blocked flows (is_blocked=1), use dpi_engine with rules:")
    print("       Windows : .\\build\\Release\\dpi_engine.exe test_dpi.pcap output.pcap --block-ip 192.168.1.50")
    print()
    print("  3. Retrain the ML model:")
    print("       python ml/train_classifier.py")


if __name__ == '__main__':
    main()