darkproxy/ensdns/server.py

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

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"))
RECONNECT_INTERVAL = float(os.getenv("ENSDNS_RECONNECT_INTERVAL", "30"))

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"


class ENSBackendUnavailable(RuntimeError):
    pass


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.rpc_url = rpc_url
        self._lock = threading.Lock()
        self._web3 = None
        self._registry = None
        self._last_connect_attempt = 0.0
        self._last_error = None

    def _connect(self, force: bool = False) -> bool:
        now = time.monotonic()
        with self._lock:
            if self._web3 is not None and self._registry is not None and self._web3.is_connected():
                return True
            if not force and now - self._last_connect_attempt < RECONNECT_INTERVAL:
                return False

            self._last_connect_attempt = now
            web3 = Web3(HTTPProvider(self.rpc_url, request_kwargs={"timeout": 10}))
            if not web3.is_connected():
                self._web3 = None
                self._registry = None
                self._last_error = f"Could not connect to Ethereum RPC: {self.rpc_url}"
                print(f"[ensdns] {self._last_error}", flush=True)
                return False

            self._web3 = web3
            self._registry = web3.eth.contract(address=ENS_REGISTRY_ADDRESS, abi=ENS_REGISTRY_ABI)
            self._last_error = None
            print(f"[ensdns] connected to {self.rpc_url}", flush=True)
            return True

    def _ensure_backend(self):
        if not self._connect():
            detail = self._last_error or f"Could not connect to Ethereum RPC: {self.rpc_url}"
            raise ENSBackendUnavailable(detail)
        return self._web3, self._registry

    def resolve(self, name: str):
        web3, registry = self._ensure_backend()
        node = namehash(name)
        try:
            resolver_addr = registry.functions.resolver(node).call()
        except Exception as exc:
            with self._lock:
                self._web3 = None
                self._registry = None
                self._last_error = f"ENS registry lookup failed: {exc}"
            raise ENSBackendUnavailable(self._last_error) from exc

        if not resolver_addr or resolver_addr == EMPTY_ADDRESS:
            return None

        resolver = 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

    try:
        resolved = ens.resolve(qname)
    except ENSBackendUnavailable as exc:
        print(f"[ensdns] backend unavailable while resolving {qname}: {exc}", flush=True)
        reply.header.rcode = RCODE.SERVFAIL
        return reply

    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] listening on {LISTEN_HOST}:{LISTEN_PORT} (udp/tcp)", flush=True)
    if not ens._connect(force=True):
        print(f"[ensdns] starting degraded; RPC unavailable at {RPC_URL}", 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()