#!/usr/bin/env python3
"""Minimal NFQUEUE example: rewrite DNS queries from example.com to pwned.com."""

import signal
import sys

from netfilterqueue import NetfilterQueue
from scapy.all import DNS, DNSQR, IP, UDP

SOURCE_QNAME = b"example.com."
TARGET_QNAME = b"pwned.com."

nfq = NetfilterQueue()


def _handle_packet(packet) -> None:
    try:
        ip = IP(packet.get_payload())
        if ip.haslayer(UDP) and ip.haslayer(DNS) and ip.haslayer(DNSQR):
            dns = ip[DNS]
            query = ip[DNSQR]

            # Only touch DNS requests for exactly example.com.
            if dns.qr == 0 and query.qname == SOURCE_QNAME:
                query.qname = TARGET_QNAME
                del ip.len
                del ip.chksum
                del ip[UDP].len
                del ip[UDP].chksum
                packet.set_payload(bytes(ip))
    except Exception:
        pass

    packet.accept()


def _stop(_sig, _frame) -> None:
    raise SystemExit(0)


def main() -> int:
    if len(sys.argv) != 2:
        print("Usage: dns_rewrite_example_to_pwned.py <qnum>", file=sys.stderr)
        return 1

    try:
        qnum = int(sys.argv[1])
    except ValueError:
        print("qnum must be an integer", file=sys.stderr)
        return 1

    signal.signal(signal.SIGINT, _stop)
    signal.signal(signal.SIGTERM, _stop)

    nfq.bind(qnum, _handle_packet)
    try:
        nfq.run()
    except KeyboardInterrupt:
        pass
    finally:
        nfq.unbind()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())