#!/usr/bin/env python3
from scapy.all import rdpcap, wrpcap, Raw, IP, IPv6, TCP, UDP
import argparse

SEARCH = b"EntySec"
REPLACE = b"CTF"

def patch_packet(pkt):
    """
    If the packet Raw payload contains 'EntySec', replace it with 'CTF'
    and clear checksums/lengths for recalculation.
    Returns (pkt, modified: bool).
    """
    if Raw not in pkt:
        return pkt, False

    raw = pkt[Raw]
    data = raw.load

    if SEARCH not in data:
        return pkt, False

    new_data = data.replace(SEARCH, REPLACE)

    # If somehow nothing changed, mark as not modified
    if new_data == data:
        return pkt, False

    raw.load = new_data

    # Clear checksums/lengths so Scapy recalculates them on write
    if IP in pkt:
        ip = pkt[IP]
        if hasattr(ip, "len"):
            del ip.len
        if hasattr(ip, "chksum"):
            del ip.chksum

    if IPv6 in pkt:
        ipv6 = pkt[IPv6]
        if hasattr(ipv6, "plen"):
            del ipv6.plen

    if TCP in pkt:
        tcp = pkt[TCP]
        if hasattr(tcp, "chksum"):
            del tcp.chksum

    if UDP in pkt:
        udp = pkt[UDP]
        if hasattr(udp, "len"):
            del udp.len
        if hasattr(udp, "chksum"):
            del udp.chksum

    return pkt, True

def main():
    parser = argparse.ArgumentParser(
        description="Replace 'EntySec' with 'flag' in PCAP payloads"
    )
    parser.add_argument("input", help="Input pcap/pcapng file")
    parser.add_argument("output", help="Output pcap/pcapng file")
    args = parser.parse_args()

    print(f"[+] Reading packets from {args.input} ...")
    pkts = rdpcap(args.input)

    new_pkts = []
    modified_count = 0

    for idx, pkt in enumerate(pkts):
        new_pkt, modified = patch_packet(pkt)
        if modified:
            modified_count += 1
            # Optional debug:
            # print(f"[+] Modified packet #{idx}")
        new_pkts.append(new_pkt)

    print(f"[+] Modified {modified_count} packet(s)")
    print(f"[+] Writing patched capture to {args.output} ...")
    wrpcap(args.output, new_pkts)
    print("[+] Done.")

if __name__ == "__main__":
    main()