#!/usr/bin/env python3
import argparse
import http.client
import queue
import socket
import ssl
import struct
import threading


OPCODE_UPDATE = 5
TYPE_A = 1
CLASS_IN = 1


def encode_name(name):
    return b"".join(bytes((len(label),)) + label.encode() for label in name.rstrip(".").split(".")) + b"\x00"


def update_message():
    zone = encode_name("example.com.") + struct.pack("!HH", 6, CLASS_IN)
    update = (
        encode_name("foo.example.com.")
        + struct.pack("!HHIH", TYPE_A, CLASS_IN, 300, 4)
        + socket.inet_aton("192.0.2.123")
    )
    header = struct.pack("!HHHHHH", 0x1234, OPCODE_UPDATE << 11, 1, 0, 1, 0)
    return header + zone + update


def read_name(message, offset):
    labels = []
    end = None
    seen = set()
    while True:
        if offset >= len(message) or offset in seen:
            raise ValueError("invalid DNS name")
        seen.add(offset)
        length = message[offset]
        if length & 0xC0 == 0xC0:
            if offset + 1 >= len(message):
                raise ValueError("truncated compression pointer")
            if end is None:
                end = offset + 2
            offset = ((length & 0x3F) << 8) | message[offset + 1]
            continue
        offset += 1
        if length == 0:
            return ".".join(labels) + ".", end if end is not None else offset
        if length & 0xC0 or offset + length > len(message):
            raise ValueError("invalid DNS label")
        labels.append(message[offset : offset + length].decode("ascii"))
        offset += length


def question_end(message, count):
    offset = 12
    for _ in range(count):
        _, offset = read_name(message, offset)
        offset += 4
        if offset > len(message):
            raise ValueError("truncated question")
    return offset


def parse_update(message):
    _, flags, qdcount, _, nscount, _ = struct.unpack_from("!HHHHHH", message)
    if (flags >> 11) & 0xF != OPCODE_UPDATE or qdcount != 1 or nscount < 1:
        return None
    offset = question_end(message, qdcount)
    name, offset = read_name(message, offset)
    rrtype, rrclass, ttl, rdlength = struct.unpack_from("!HHIH", message, offset)
    offset += 10
    rdata = message[offset : offset + rdlength]
    if rrtype != TYPE_A or rrclass != CLASS_IN or len(rdata) != 4:
        return None
    return name, ttl, socket.inet_ntoa(rdata)


def response_for(message):
    ident, flags, qdcount, _, _, _ = struct.unpack_from("!HHHHHH", message)
    end = question_end(message, qdcount)
    response_flags = flags | 0x8000
    return struct.pack("!HHHHHH", ident, response_flags, qdcount, 0, 0, 0) + message[12:end]


def handle_message(message, peer, received):
    try:
        update = parse_update(message)
        response = response_for(message)
    except (ValueError, struct.error):
        return None
    if update is not None:
        received.put((peer, update))
    return response


def serve_udp(sock, received, stopped):
    while not stopped.is_set():
        try:
            message, peer = sock.recvfrom(65535)
        except socket.timeout:
            continue
        except OSError:
            return
        response = handle_message(message, peer, received)
        if response is not None:
            sock.sendto(response, peer)


def recv_exact(connection, size, stopped):
    data = bytearray()
    while len(data) < size and not stopped.is_set():
        try:
            chunk = connection.recv(size - len(data))
        except socket.timeout:
            continue
        if not chunk:
            return None
        data.extend(chunk)
    return bytes(data) if len(data) == size else None


def serve_tcp(sock, received, stopped):
    while not stopped.is_set():
        try:
            connection, peer = sock.accept()
        except socket.timeout:
            continue
        except OSError:
            return
        with connection:
            connection.settimeout(0.1)
            while not stopped.is_set():
                length = recv_exact(connection, 2, stopped)
                if length is None:
                    break
                message = recv_exact(connection, struct.unpack("!H", length)[0], stopped)
                if message is None:
                    break
                response = handle_message(message, peer, received)
                if response is not None:
                    connection.sendall(struct.pack("!H", len(response)) + response)


def send_doh(host, port, payload, timeout):
    context = ssl._create_unverified_context()
    connection = http.client.HTTPSConnection(host, port, timeout=timeout, context=context)
    try:
        connection.request(
            "POST",
            "/dns-query",
            body=payload,
            headers={"Content-Type": "application/dns-message"},
        )
        response = connection.getresponse()
        body = response.read()
        return response.status, len(body)
    finally:
        connection.close()


def main():
    parser = argparse.ArgumentParser(description="Probe whether CoreDNS forwards RFC 2136 UPDATE over DoH")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8053)
    parser.add_argument("--upstream-host", default="127.0.0.1")
    parser.add_argument("--upstream-port", type=int, default=15354)
    parser.add_argument("--timeout", type=float, default=2.0)
    parser.add_argument("--expect", choices=("forwarded", "blocked", "either"), default="either")
    args = parser.parse_args()

    received = queue.Queue()
    stopped = threading.Event()
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.settimeout(0.1)
    udp_sock.bind((args.upstream_host, args.upstream_port))
    tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_sock.settimeout(0.1)
    tcp_sock.bind((args.upstream_host, args.upstream_port))
    tcp_sock.listen()
    threads = [
        threading.Thread(target=serve_udp, args=(udp_sock, received, stopped), daemon=True),
        threading.Thread(target=serve_tcp, args=(tcp_sock, received, stopped), daemon=True),
    ]
    for thread in threads:
        thread.start()

    payload = update_message()
    try:
        status, response_bytes = send_doh(args.host, args.port, payload, args.timeout)
        try:
            peer, update = received.get(timeout=args.timeout)
        except queue.Empty:
            peer = update = None
    finally:
        stopped.set()
        udp_sock.close()
        tcp_sock.close()
        for thread in threads:
            thread.join(timeout=1)

    print("payload=%d opcode=UPDATE record=foo.example.com. 300 IN A 192.0.2.123" % len(payload))
    print("http_status=%d response_bytes=%d" % (status, response_bytes))
    if update is None:
        result = "blocked"
        print("upstream_received_update=false")
    else:
        result = "forwarded"
        name, ttl, address = update
        print("upstream_received_update=true source=%s:%d" % peer)
        print("upstream_record=%s %d IN A %s" % (name, ttl, address))
    print("result=%s" % result)

    if args.expect != "either" and args.expect != result:
        raise SystemExit("expected %s, got %s" % (args.expect, result))


if __name__ == "__main__":
    main()