TCP reset connection tool

rst_inject.py

#!/usr/bin/env python3
"""
rst_inject.py - Inject TCP RST packets into an existing TCP stream to
reproduce application bugs related to abrupt connection termination.
"""

import argparse
import sys
import time
from dataclasses import dataclass
from typing import Optional

from scapy.all import IP, TCP, L3RawSocket, conf, sniff


@dataclass
class Flow:
    src: str
    sport: int
    dst: str
    dport: int
    last_seq: int = 0
    last_ack: int = 0
    seen: bool = False


def build_rst(flow: Flow, direction: str = "to_dst", ack_flag: bool = True) -> IP:
    if direction == "to_dst":
        ip_layer = IP(src=flow.src, dst=flow.dst)
        sport, dport = flow.sport, flow.dport
        seq = flow.last_seq
        ack = flow.last_ack
    else:
        ip_layer = IP(src=flow.dst, dst=flow.src)
        sport, dport = flow.dport, flow.sport
        seq = flow.last_ack
        ack = flow.last_seq

    flags = "RA" if ack_flag else "R"
    tcp_layer = TCP(sport=sport, dport=dport, flags=flags, seq=seq, ack=ack, window=0)
    return ip_layer / tcp_layer


def get_l3_socket(iface: Optional[str]):
    """
    Build an L3 socket explicitly bound to `iface`, if given.
    Falls back to scapy's default L3 socket (routed via table) otherwise.
    """
    if iface:
        # conf.L3socket is the platform-appropriate class (e.g. L3RawSocket on macOS/BSD)
        return conf.L3socket(iface=iface)
    return conf.L3socket()


def send_rst(
    flow: Flow,
    direction: str = "to_dst",
    count: int = 1,
    iface: Optional[str] = None,
    verbose: bool = True,
):
    pkt = build_rst(flow, direction=direction)
    sock = get_l3_socket(iface)
    try:
        for _ in range(count):
            sock.send(pkt)
    finally:
        sock.close()
    if verbose:
        print(
            f"[+] Sent {count} RST packet(s) direction={direction} "
            f"seq={pkt[TCP].seq} ack={pkt[TCP].ack} "
            f"{pkt[IP].src}:{pkt[TCP].sport} -> {pkt[IP].dst}:{pkt[TCP].dport}"
            f"{' via ' + iface if iface else ''}"
        )


def sniff_and_inject(
    host: str,
    port: int,
    iface: Optional[str],
    direction: str,
    delay: float,
    count: int,
    only_first: bool,
):
    flow = Flow(src="", sport=0, dst=host, dport=port)
    injected = {"done": False}

    bpf_filter = f"tcp and host {host} and port {port}"

    def handle(pkt):
        if injected["done"] and only_first:
            return
        if IP not in pkt or TCP not in pkt:
            return

        ip = pkt[IP]
        tcp = pkt[TCP]

        if ip.dst == host and tcp.dport == port:
            flow.src, flow.sport = ip.src, tcp.sport
            flow.dst, flow.dport = ip.dst, tcp.dport
            flow.last_seq = tcp.seq + max(len(tcp.payload), 1)
            flow.seen = True
        elif ip.src == host and tcp.sport == port:
            flow.last_ack = tcp.seq + max(len(tcp.payload), 1)
            flow.seen = True

        if flow.seen and flow.last_seq and flow.last_ack:
            if delay:
                print(f"[*] Stream captured, waiting {delay}s before injecting RST...")
                time.sleep(delay)
            send_rst(flow, direction=direction, count=count, iface=iface)
            injected["done"] = True
            if only_first:
                raise KeyboardInterrupt

    print(f"[*] Sniffing on iface={iface or conf.iface} for {host}:{port} ...")
    try:
        sniff(filter=bpf_filter, prn=handle, iface=iface, store=False)
    except KeyboardInterrupt:
        pass


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--iface",
        help="Network interface (e.g. en0). Run `ifconfig` to list options on macOS.",
    )
    parser.add_argument("--direction", choices=["to_dst", "to_src"], default="to_dst")
    parser.add_argument("--count", type=int, default=1)

    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--sniff", action="store_true")
    mode.add_argument("--direct", action="store_true")

    parser.add_argument("--host")
    parser.add_argument("--port", type=int)
    parser.add_argument("--delay", type=float, default=0.0)
    parser.add_argument("--only-first", action="store_true", default=True)

    parser.add_argument("--src")
    parser.add_argument("--sport", type=int)
    parser.add_argument("--dst")
    parser.add_argument("--dport", type=int)
    parser.add_argument("--seq", type=int)
    parser.add_argument("--ack", type=int, default=0)

    args = parser.parse_args()

    if args.sniff:
        if not args.host or not args.port:
            parser.error("--sniff requires --host and --port")
        sniff_and_inject(
            host=args.host,
            port=args.port,
            iface=args.iface,
            direction=args.direction,
            delay=args.delay,
            count=args.count,
            only_first=args.only_first,
        )
    else:
        required = [args.src, args.sport, args.dst, args.dport, args.seq]
        if any(v is None for v in required):
            parser.error("--direct requires --src --sport --dst --dport --seq")
        flow = Flow(
            src=args.src,
            sport=args.sport,
            dst=args.dst,
            dport=args.dport,
            last_seq=args.seq,
            last_ack=args.ack,
        )
        send_rst(flow, direction=args.direction, count=args.count, iface=args.iface)


if __name__ == "__main__":
    if sys.platform != "darwin":
        print(
            "[!] Warning: this script was tailored for macOS but should run on "
            "other platforms too.",
            file=sys.stderr,
        )
    main()