#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RRPP UoT v2 relay — UDP over TCP 解封装服务（sing-box UoT v2 协议）

协议（SagerNet 专有，见 sing-box 文档 configuration/shared/udp-over-tcp）：

  请求头:  | isConnect u8 | ATYP u8 | address | port u16be |
    isConnect=1: 固定目标模式，后续帧全部是 length+data，发往请求头里的目标
    isConnect=0: 每帧自带目标，帧格式同 v1: | ATYP | address | port | length u16be | data |

  connect 帧:      | length u16be | data |
  non-connect 帧:  | ATYP | address | port | length u16be | data |

监听 127.0.0.1，仅供本机 nginx proxy_connect 转发，不对公网开放。
"""
import socket
import struct
import threading
import select
import time
import sys

LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 15201
IDLE_TIMEOUT = 540          # 9 分钟无流量断开（须 < nginx proxy_connect_data_timeout 10m）
MAX_DGRAM = 65507           # UDP over IPv4 单包上限
LOG_PREFIX = "[uot-relay]"


def log(msg):
    print(f"{LOG_PREFIX} {time.strftime('%Y-%m-%d %H:%M:%S')} {msg}", flush=True)


def read_exact(f, n):
    buf = b""
    while len(buf) < n:
        chunk = f.read(n - len(buf))
        if not chunk:
            raise EOFError("connection closed")
        buf += chunk
    return buf


def parse_addr(f):
    """从流中读取 SOCKS 格式地址（ATYP + address + port）。"""
    atyp = read_exact(f, 1)[0]
    if atyp == 0x00:
        addr = socket.inet_ntop(socket.AF_INET, read_exact(f, 4))
    elif atyp == 0x01:
        addr = socket.inet_ntop(socket.AF_INET6, read_exact(f, 16))
    elif atyp == 0x02:
        ln = read_exact(f, 1)[0]
        addr = read_exact(f, ln).decode("utf-8", "replace")
    else:
        raise ValueError(f"bad ATYP {atyp:#x}")
    port = struct.unpack(">H", read_exact(f, 2))[0]
    return addr, port


def encode_addr(addr, port):
    """SOCKS 格式地址编码（用于 non-connect 模式的回包）。"""
    try:
        return b"\x00" + socket.inet_pton(socket.AF_INET, addr) + struct.pack(">H", port)
    except OSError:
        pass
    try:
        return b"\x01" + socket.inet_pton(socket.AF_INET6, addr) + struct.pack(">H", port)
    except OSError:
        pass
    d = addr.encode("utf-8")
    if len(d) > 255:
        raise ValueError("domain too long")
    return b"\x02" + bytes([len(d)]) + d + struct.pack(">H", port)


def handle(conn, peer):
    conn.settimeout(IDLE_TIMEOUT)
    f = conn.makefile("rb")
    udp = None
    stop = threading.Event()
    try:
        is_connect = read_exact(f, 1)[0]
        req_dst = parse_addr(f)
        log(f"conn from nginx: isConnect={is_connect} dst={req_dst}")

        udp = socket.socket(socket.AF_INET6 if ":" in req_dst[0] else socket.AF_INET,
                            socket.SOCK_DGRAM)
        udp.setblocking(False)
        if is_connect:
            udp.connect(req_dst)  # 固定目标：连通式 UDP

        write_lock = threading.Lock()

        def tcp_write(data):
            with write_lock:
                conn.sendall(data)

        def uplink():
            """TCP -> UDP：按格式读帧并发送。"""
            try:
                while not stop.is_set():
                    if is_connect:
                        ln = struct.unpack(">H", read_exact(f, 2))[0]
                        if ln > MAX_DGRAM:
                            raise ValueError(f"datagram too large: {ln}")
                        data = read_exact(f, ln)
                        udp.send(data)
                    else:
                        dst = parse_addr(f)
                        ln = struct.unpack(">H", read_exact(f, 2))[0]
                        if ln > MAX_DGRAM:
                            raise ValueError(f"datagram too large: {ln}")
                        data = read_exact(f, ln)
                        udp.sendto(data, dst)
            except Exception as e:
                if not stop.is_set():
                    log(f"uplink end: {type(e).__name__}: {e}")
            finally:
                stop.set()

        t = threading.Thread(target=uplink, daemon=True)
        t.start()

        # 主循环：UDP -> TCP，带空闲超时
        last = time.time()
        while not stop.is_set():
            r, _, _ = select.select([udp], [], [], 2.0)
            if r:
                if is_connect:
                    data = udp.recv(MAX_DGRAM)
                    tcp_write(struct.pack(">H", len(data)) + data)
                else:
                    data, src = udp.recvfrom(MAX_DGRAM)
                    tcp_write(encode_addr(*src) + struct.pack(">H", len(data)) + data)
                last = time.time()
            elif time.time() - last > IDLE_TIMEOUT:
                log("idle timeout, closing")
                break
    except Exception as e:
        log(f"conn error: {type(e).__name__}: {e}")
    finally:
        stop.set()
        try:
            f.close()
        except Exception:
            pass
        try:
            conn.close()
        except Exception:
            pass
        if udp:
            try:
                udp.close()
            except Exception:
                pass


def main():
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((LISTEN_HOST, LISTEN_PORT))
    srv.listen(128)
    log(f"listening on {LISTEN_HOST}:{LISTEN_PORT}")
    while True:
        conn, peer = srv.accept()
        threading.Thread(target=handle, args=(conn, peer), daemon=True).start()


if __name__ == "__main__":
    main()