Remove payload from a pcap (useful to fully anonymize a pcap)

pcap-remove-payload.py

#! /usr/bin/env python2
from scapy.all import *
import sys


INFILE = sys.argv[1]
OUTFILE = sys.argv[2]

with PcapWriter(OUTFILE) as dest:
    with PcapReader(INFILE) as infile:
        for pkt in infile:
            if TCP in pkt:
                pkt[TCP].remove_payload()
            elif UDP in pkt:
                pkt[UDP].remove_payload()
            dest.write(pkt)

I apply it in python3. But it cannot reference to TCP and UDP. What's wrong?Could you give me some advice?

This should work (python3):

#! /usr/bin/env python3
from scapy.utils import rdpcap, wrpcap
from scapy.layers.inet import IP, TCP, UDP  # import needed!
import sys

INFILE = sys.argv[1]
OUTFILE = sys.argv[2]

packets = rdpcap(INFILE)

new_packets = []
for packet in packets:
    if packet.haslayer("IP"):
        if packet.haslayer("TCP"):
            packet["IP"]["TCP"].remove_payload()
        elif packet.haslayer("UDP"):
            packet["IP"]["UDP"].remove_payload()

    new_packets.append(packet)

wrpcap(OUTFILE, new_packets)

Would it be possible to do such postprocessing on the fly?

#!/usr/bin/env python3
from scapy.all import *
from scapy.layers.inet import *
import sys

with PcapWriter(sys.argv[2]) as out:
    with PcapReader(sys.argv[1]) as in:
        for pkt in in:
            if   TCP in pkt: pkt[TCP].remove_payload()
            elif UDP in pkt: pkt[UDP].remove_payload()
            out.write(pkt)