# WinDivert DNS responder for XTOS readers that hard-code 223.5.5.5 (crosspoint-tools#51).
# Answers A queries for the Xteink update hosts from hotspot clients with the Unlocker's address,
# passes every other packet through untouched. Run as Administrator; Ctrl+C to stop.
import ctypes
import os
import random
import struct
import sys
import threading
import time
from ctypes import wintypes

DLL = r"C:\Program Files\Xteink Unlocker\WinDivert.dll"
SPOOF_HOSTS = {"api-prod.xteink.cn", "api-prod.xteink.cc"}
SPOOF_IP = "192.168.137.1"
FILTER = (b"ip and udp.DstPort == 53 and ip.SrcAddr >= 192.168.137.2"
          b" and ip.SrcAddr <= 192.168.137.254")

LAYER_NETWORK, LAYER_NETWORK_FORWARD = 0, 1
FLAG_SEND_ONLY = 0x0008
ADDR_SIZE = 80  # sizeof(WINDIVERT_ADDRESS) in WinDivert 2.x
OUTBOUND_BIT = 1 << 17
INVALID_HANDLE = ctypes.c_void_p(-1).value

wd = ctypes.WinDLL(DLL, use_last_error=True)
wd.WinDivertOpen.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_int16, ctypes.c_uint64]
wd.WinDivertOpen.restype = ctypes.c_void_p
wd.WinDivertRecv.argtypes = [ctypes.c_void_p, ctypes.c_void_p, wintypes.UINT,
                             ctypes.POINTER(wintypes.UINT), ctypes.c_void_p]
wd.WinDivertRecv.restype = wintypes.BOOL
wd.WinDivertSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, wintypes.UINT,
                             ctypes.POINTER(wintypes.UINT), ctypes.c_void_p]
wd.WinDivertSend.restype = wintypes.BOOL
wd.WinDivertHelperCalcChecksums.argtypes = [ctypes.c_void_p, wintypes.UINT, ctypes.c_void_p, ctypes.c_uint64]
wd.WinDivertHelperCalcChecksums.restype = wintypes.BOOL


def open_handle(layer, filt, flags=0):
    h = wd.WinDivertOpen(filt, layer, 0, flags)
    if h is None or h == INVALID_HANDLE:
        raise ctypes.WinError(ctypes.get_last_error())
    return h


def log(msg):
    print(time.strftime("%H:%M:%S"), msg, flush=True)


def parse_query(pkt):
    """Return (ihl, qname, qtype, question_end) for an IPv4/UDP DNS query, or None."""
    ihl = (pkt[0] & 0x0F) * 4
    dns = ihl + 8
    if len(pkt) < dns + 12 or pkt[dns + 2] & 0x80:  # too short, or a response
        return None
    labels, i = [], dns + 12
    while i < len(pkt) and pkt[i]:
        if pkt[i] & 0xC0:
            return None
        labels.append(pkt[i + 1:i + 1 + pkt[i]].decode(errors="replace"))
        i += 1 + pkt[i]
    if i + 5 > len(pkt):
        return None
    qtype = struct.unpack_from("!H", pkt, i + 1)[0]
    return ihl, ".".join(labels).lower(), qtype, i + 5


def build_reply(pkt, ihl, qtype, qend):
    dns = ihl + 8
    qid, qflags = struct.unpack_from("!HH", pkt, dns)
    question = pkt[dns + 12:qend]
    answers = b""
    if qtype == 1:  # A -> Unlocker; anything else (AAAA, HTTPS...) -> empty NOERROR
        answers = struct.pack("!HHHIH", 0xC00C, 1, 1, 60, 4) + bytes(map(int, SPOOF_IP.split(".")))
    flags = 0x8080 | (qflags & 0x0100)  # QR + RA, copy RD
    payload = struct.pack("!HHHHHH", qid, flags, 1, 1 if answers else 0, 0, 0) + question + answers
    src_ip, dst_ip = pkt[12:16], pkt[16:20]
    sport, = struct.unpack_from("!H", pkt, ihl)
    udp = struct.pack("!HHHH", 53, sport, 8 + len(payload), 0) + payload
    ip = struct.pack("!BBHHHBBH4s4s", 0x45, 0, 20 + len(udp), random.randrange(65536), 0,
                     64, 17, 0, dst_ip, src_ip)
    return ip + udp


def worker(name, handle, sender):
    buf = ctypes.create_string_buffer(65535)
    addr = ctypes.create_string_buffer(ADDR_SIZE)
    n = wintypes.UINT()
    while True:
        if not wd.WinDivertRecv(handle, buf, len(buf), ctypes.byref(n), addr):
            log(f"[{name}] recv error {ctypes.get_last_error()}")
            continue
        pkt = buf.raw[:n.value]
        q = parse_query(pkt)
        src = ".".join(map(str, pkt[12:16]))
        dst = ".".join(map(str, pkt[16:20]))
        if q and q[1] in SPOOF_HOSTS:
            ihl, qname, qtype, qend = q
            reply = ctypes.create_string_buffer(build_reply(pkt, ihl, qtype, qend))
            rlen = len(reply) - 1  # drop create_string_buffer's trailing NUL
            raddr = ctypes.create_string_buffer(ADDR_SIZE)
            bits = LAYER_NETWORK | OUTBOUND_BIT
            struct.pack_into("<I", raddr, 8, bits)
            wd.WinDivertHelperCalcChecksums(reply, rlen, raddr, 0)
            ok = wd.WinDivertSend(sender, reply, rlen, None, raddr)
            log(f"[{name}] SPOOF {src} -> {dst}  {qname} type {qtype} -> {SPOOF_IP}"
                + ("" if ok else f"  (send error {ctypes.get_last_error()})"))
        else:
            wd.WinDivertSend(handle, buf, n.value, None, addr)
            log(f"[{name}] pass  {src} -> {dst}  {q[1] if q else '?'}")


def main():
    if not ctypes.windll.shell32.IsUserAnAdmin():
        sys.exit("Run as Administrator.")
    if not os.path.exists(DLL):
        sys.exit(f"{DLL} not found.")
    sender = open_handle(LAYER_NETWORK, b"false", FLAG_SEND_ONLY)
    for name, layer in (("network", LAYER_NETWORK), ("forward", LAYER_NETWORK_FORWARD)):
        h = open_handle(layer, FILTER)
        threading.Thread(target=worker, args=(name, h, sender), daemon=True).start()
    log(f"diverting DNS from hotspot clients; spoofing {', '.join(sorted(SPOOF_HOSTS))} -> {SPOOF_IP}")
    log("press 'Check for Updates' on the reader now (Ctrl+C to stop)")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        log("stopped")


if __name__ == "__main__":
    main()