#!/usr/bin/env python3
"""RRPP 裸 UoT v2 relay。

nginx 的 proxy_connect 模块负责 TLS、Basic Auth 和 CONNECT，relay 只监听
127.0.0.1:15201，收到的首字节就是 UoT v2 request，不再解析 HTTP。
"""

from __future__ import annotations

import ipaddress
import logging
import socket
import struct
import threading

MAGIC = "sp.v2.udp-over-tcp.arpa"
LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 15201
MAX_PACKET = 65535
IDLE_TIMEOUT = 9 * 60


def read_exact(conn: socket.socket, size: int) -> bytes:
    data = bytearray()
    while len(data) < size:
        chunk = conn.recv(size - len(data))
        if not chunk:
            raise ConnectionError("客户端连接已关闭")
        data.extend(chunk)
    return bytes(data)


def read_address(conn: socket.socket):
    atyp = read_exact(conn, 1)[0]
    if atyp == 0:
        host = str(ipaddress.IPv4Address(read_exact(conn, 4)))
    elif atyp == 1:
        host = str(ipaddress.IPv6Address(read_exact(conn, 16)))
    elif atyp == 2:
        size = read_exact(conn, 1)[0]
        host = read_exact(conn, size).decode("idna")
    else:
        raise ValueError(f"未知 UoT 地址类型: {atyp}")
    return host, struct.unpack("!H", read_exact(conn, 2))[0]


def encode_address(host: str, port: int) -> bytes:
    try:
        addr = ipaddress.ip_address(host)
    except ValueError:
        raw = host.encode("idna")
        return b"\x02" + bytes([len(raw)]) + raw + struct.pack("!H", port)
    return (b"\x00" if addr.version == 4 else b"\x01") + addr.packed + struct.pack("!H", port)


def resolve_destination(host: str, port: int):
    return socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_DGRAM)[0][4]


def handle(conn: socket.socket, peer):
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.settimeout(0.2)
    # 与 nginx 的 10 分钟读写超时错开，空闲连接由 relay 先回收。
    conn.settimeout(IDLE_TIMEOUT)
    send_lock = threading.Lock()
    try:
        is_connect = read_exact(conn, 1)[0]
        first_target = read_address(conn)
        logging.info("UoT client %s target=%s:%s connect=%s", peer, *first_target, is_connect)

        def udp_reader():
            while True:
                try:
                    data, source = udp.recvfrom(MAX_PACKET)
                except socket.timeout:
                    continue
                except OSError:
                    return
                try:
                    with send_lock:
                        if is_connect:
                            conn.sendall(struct.pack("!H", len(data)) + data)
                        else:
                            conn.sendall(encode_address(source[0], source[1]) +
                                          struct.pack("!H", len(data)) + data)
                except OSError:
                    return

        threading.Thread(target=udp_reader, daemon=True).start()
        if is_connect:
            while True:
                length = struct.unpack("!H", read_exact(conn, 2))[0]
                payload = read_exact(conn, length)
                udp.sendto(payload, resolve_destination(*first_target))
        else:
            while True:
                destination = read_address(conn)
                length = struct.unpack("!H", read_exact(conn, 2))[0]
                payload = read_exact(conn, length)
                udp.sendto(payload, resolve_destination(*destination))
    except (ConnectionError, OSError, ValueError) as exc:
        logging.info("UoT client %s closed: %s", peer, exc)
    finally:
        udp.close()
        conn.close()


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind((LISTEN_HOST, LISTEN_PORT))
    listener.listen(128)
    logging.info("RRPP UoT relay listening on %s:%d", LISTEN_HOST, LISTEN_PORT)
    while True:
        conn, peer = listener.accept()
        threading.Thread(target=handle, args=(conn, peer), daemon=True).start()


if __name__ == "__main__":
    main()