#!/usr/bin/env python3
import os
import socket
import threading

from dnslib import A, AAAA, CNAME, QTYPE, RR, TXT, DNSHeader, DNSRecord, RCODE
from web3 import HTTPProvider, Web3


RPC_URL = os.getenv("ENS_RPC_URL", "https://ethereum-rpc.publicnode.com")
LISTEN_HOST = os.getenv("ENSDNS_LISTEN_HOST", "0.0.0.0")
LISTEN_PORT = int(os.getenv("ENSDNS_LISTEN_PORT", "53"))
DEFAULT_TTL = int(os.getenv("ENSDNS_TTL", "60"))

ENS_REGISTRY_ADDRESS = Web3.to_checksum_address("0x00000000000C2E074eC69A0dFb2997BA6C7d2e1e")

ENS_REGISTRY_ABI = [
    {
        "constant": True,
        "inputs": [{"name": "node", "type": "bytes32"}],
        "name": "resolver",
        "outputs": [{"name": "", "type": "address"}],
        "type": "function",
    }
]

PUBLIC_RESOLVER_ABI = [
    {
        "constant": True,
        "inputs": [{"name": "node", "type": "bytes32"}],
        "name": "addr",
        "outputs": [{"name": "", "type": "address"}],
        "type": "function",
    },
    {
        "constant": True,
        "inputs": [{"name": "node", "type": "bytes32"}],
        "name": "contenthash",
        "outputs": [{"name": "", "type": "bytes"}],
        "type": "function",
    },
]

EMPTY_ADDRESS = "0x0000000000000000000000000000000000000000"


def namehash(name: str) -> bytes:
    node = b"\x00" * 32
    labels = [label for label in name.strip().lower().split(".") if label]
    for label in reversed(labels):
        label_hash = Web3.keccak(text=label)
        node = Web3.keccak(node + label_hash)
    return node


def normalize_qname(qname: str) -> str:
    return qname.rstrip(".").lower()


class ENSResolver:
    def __init__(self, rpc_url: str):
        self.web3 = Web3(HTTPProvider(rpc_url, request_kwargs={"timeout": 10}))
        if not self.web3.is_connected():
            raise RuntimeError(f"Could not connect to Ethereum RPC: {rpc_url}")
        self.registry = self.web3.eth.contract(address=ENS_REGISTRY_ADDRESS, abi=ENS_REGISTRY_ABI)

    def resolve(self, name: str):
        node = namehash(name)
        resolver_addr = self.registry.functions.resolver(node).call()
        if not resolver_addr or resolver_addr == EMPTY_ADDRESS:
            return None

        resolver = self.web3.eth.contract(address=resolver_addr, abi=PUBLIC_RESOLVER_ABI)

        addr = None
        contenthash = None

        try:
            raw_addr = resolver.functions.addr(node).call()
            if raw_addr and raw_addr != EMPTY_ADDRESS:
                addr = Web3.to_checksum_address(raw_addr)
        except Exception:
            addr = None

        try:
            raw_ch = resolver.functions.contenthash(node).call()
            if raw_ch:
                contenthash = "0x" + raw_ch.hex()
        except Exception:
            contenthash = None

        return {
            "address": addr,
            "contenthash": contenthash,
        }


def answer_query(record: DNSRecord, ens: ENSResolver) -> DNSRecord:
    reply = DNSRecord(DNSHeader(id=record.header.id, qr=1, aa=1, ra=1), q=record.q)

    qname = normalize_qname(str(record.q.qname))
    qtype = QTYPE[record.q.qtype]

    if not qname.endswith(".eth"):
        reply.header.rcode = RCODE.NXDOMAIN
        return reply

    resolved = ens.resolve(qname)
    if resolved is None:
        reply.header.rcode = RCODE.NXDOMAIN
        return reply

    fqdn = str(record.q.qname)

    if qtype in ("TXT", "ANY"):
        if resolved.get("address"):
            reply.add_answer(RR(fqdn, QTYPE.TXT, rdata=TXT(f"address={resolved['address']}"), ttl=DEFAULT_TTL))
        if resolved.get("contenthash"):
            reply.add_answer(RR(fqdn, QTYPE.TXT, rdata=TXT(f"contenthash={resolved['contenthash']}"), ttl=DEFAULT_TTL))

    if qtype in ("A", "AAAA", "CNAME", "ANY"):
        # Provide a browser-friendly fallback target for .eth domains.
        # e.g. vitalik.eth -> vitalik.eth.limo
        cname_target = qname + ".limo."
        reply.add_answer(RR(fqdn, QTYPE.CNAME, rdata=CNAME(cname_target), ttl=DEFAULT_TTL))

    return reply


def serve_udp(sock: socket.socket, ens: ENSResolver):
    while True:
        data, addr = sock.recvfrom(4096)
        try:
            request = DNSRecord.parse(data)
            response = answer_query(request, ens)
            sock.sendto(response.pack(), addr)
        except Exception:
            pass


def handle_tcp_conn(conn: socket.socket, ens: ENSResolver):
    try:
        length_data = conn.recv(2)
        if len(length_data) != 2:
            return
        length = int.from_bytes(length_data, "big")
        payload = conn.recv(length)
        if len(payload) != length:
            return

        request = DNSRecord.parse(payload)
        response = answer_query(request, ens).pack()
        conn.sendall(len(response).to_bytes(2, "big") + response)
    finally:
        conn.close()


def serve_tcp(sock: socket.socket, ens: ENSResolver):
    while True:
        conn, _ = sock.accept()
        threading.Thread(target=handle_tcp_conn, args=(conn, ens), daemon=True).start()


def main():
    ens = ENSResolver(RPC_URL)
    print(f"[ensdns] connected to {RPC_URL}", flush=True)
    print(f"[ensdns] listening on {LISTEN_HOST}:{LISTEN_PORT} (udp/tcp)", flush=True)

    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.bind((LISTEN_HOST, LISTEN_PORT))

    tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_sock.bind((LISTEN_HOST, LISTEN_PORT))
    tcp_sock.listen(128)

    threading.Thread(target=serve_udp, args=(udp_sock, ens), daemon=True).start()
    serve_tcp(tcp_sock, ens)


if __name__ == "__main__":
    main()