#!/usr/bin/env python3
import json
import os
import socket
import threading
from urllib.request import Request, urlopen

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


LISTEN_HOST = os.getenv("NAMECOIN_DNS_LISTEN_HOST", "0.0.0.0")
LISTEN_PORT = int(os.getenv("NAMECOIN_DNS_LISTEN_PORT", "53"))
RPC_TIMEOUT = float(os.getenv("NAMECOIN_RPC_TIMEOUT", "8"))

RPC_URL = os.getenv("NAMECOIN_RPC_URL", "http://darknamecoind:8336/")
RPC_USER = os.getenv("NAMECOIN_RPC_USER", "namecoinrpc")
RPC_PASSWORD = os.getenv("NAMECOIN_RPC_PASSWORD", "CHANGE_ME_NAMECOIN_RPC_PASSWORD")


def is_bit_query(record: DNSRecord) -> bool:
    qname = str(record.q.qname).rstrip(".").lower()
    return qname.endswith(".bit")


def rpc_call(method: str, params: list):
    payload = json.dumps({"jsonrpc": "1.0", "id": "namecoindns", "method": method, "params": params}).encode("utf-8")
    req = Request(RPC_URL, data=payload, headers={"Content-Type": "application/json"})
    token = (f"{RPC_USER}:{RPC_PASSWORD}").encode("utf-8")
    import base64

    req.add_header("Authorization", "Basic " + base64.b64encode(token).decode("ascii"))
    with urlopen(req, timeout=RPC_TIMEOUT) as resp:
        body = resp.read()
    result = json.loads(body.decode("utf-8"))
    if result.get("error"):
        raise RuntimeError(str(result["error"]))
    return result.get("result")


def lookup_namecoin_bit(label: str):
    key = f"d/{label}"
    try:
        name_info = rpc_call("name_show", [key])
    except Exception:
        return None

    value = name_info.get("value")
    if not value:
        return None

    try:
        parsed = json.loads(value)
    except Exception:
        return None

    if isinstance(parsed, dict) and "map" in parsed and isinstance(parsed["map"], dict):
        parsed = parsed["map"].get("", parsed)

    if not isinstance(parsed, dict):
        return None

    return parsed


def fail_response(payload: bytes) -> bytes:
    req = DNSRecord.parse(payload)
    reply = DNSRecord(DNSHeader(id=req.header.id, qr=1, ra=1, aa=0, rcode=RCODE.SERVFAIL), q=req.q)
    return reply.pack()


def handle_packet(payload: bytes) -> bytes:
    try:
        req = DNSRecord.parse(payload)
    except Exception:
        return b""

    if not is_bit_query(req):
        reply = DNSRecord(DNSHeader(id=req.header.id, qr=1, ra=1, aa=0, rcode=RCODE.NXDOMAIN), q=req.q)
        return reply.pack()

    qname = str(req.q.qname).rstrip(".").lower()
    label = qname[: -len(".bit")].strip(".")
    if not label:
        return fail_response(payload)

    record = lookup_namecoin_bit(label)
    if record is None:
        reply = DNSRecord(DNSHeader(id=req.header.id, qr=1, ra=1, aa=1, rcode=RCODE.NXDOMAIN), q=req.q)
        return reply.pack()

    reply = DNSRecord(DNSHeader(id=req.header.id, qr=1, ra=1, aa=1, rcode=RCODE.NOERROR), q=req.q)
    fqdn = str(req.q.qname)
    qtype = QTYPE[req.q.qtype]
    ttl = 60

    ip = record.get("ip")
    ips = ip if isinstance(ip, list) else ([ip] if isinstance(ip, str) else [])
    for item in ips:
        if ":" in item:
            if qtype in ("AAAA", "ANY"):
                reply.add_answer(RR(fqdn, QTYPE.AAAA, rdata=AAAA(item), ttl=ttl))
        else:
            if qtype in ("A", "ANY"):
                reply.add_answer(RR(fqdn, QTYPE.A, rdata=A(item), ttl=ttl))

    alias = record.get("alias") or record.get("translate")
    if isinstance(alias, str) and alias:
        target = alias if alias.endswith(".") else alias + "."
        if qtype in ("CNAME", "A", "AAAA", "ANY"):
            reply.add_answer(RR(fqdn, QTYPE.CNAME, rdata=CNAME(target), ttl=ttl))

    if qtype in ("TXT", "ANY"):
        reply.add_answer(RR(fqdn, QTYPE.TXT, rdata=TXT(json.dumps(record, separators=(",", ":"))), ttl=ttl))

    if len(reply.rr) == 0:
        return fail_response(payload)

    return reply.pack()


def serve_udp(sock: socket.socket):
    while True:
        data, addr = sock.recvfrom(4096)
        out = handle_packet(data)
        if out:
            sock.sendto(out, addr)


def handle_tcp_conn(conn: socket.socket):
    try:
        header = conn.recv(2)
        if len(header) != 2:
            return
        length = int.from_bytes(header, "big")
        payload = conn.recv(length)
        if len(payload) != length:
            return
        out = handle_packet(payload)
        conn.sendall(len(out).to_bytes(2, "big") + out)
    finally:
        conn.close()


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


def main():
    print(f"[namecoindns] listening on {LISTEN_HOST}:{LISTEN_PORT}", flush=True)
    print(f"[namecoindns] RPC endpoint: {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,), daemon=True).start()
    serve_tcp(tcp_sock)


if __name__ == "__main__":
    main()