#!/bin/bash
set -e

# ============================================================
# setup_warp_zerotrust.sh
# IPv6-only VPS 用纯 WireGuard 接入 Cloudflare WARP(不装官方客户端)
#
# 模式一(默认): 免费版 WARP —— 匿名注册,获得 IPv4 出口
#   bash setup_warp_zerotrust.sh
#
# 模式二: Zero Trust 团队设备 —— 设备显示在 Zero Trust 仪表盘,
#   受 Gateway 策略管控。仪表盘一次性准备(官方文档: Deploy the
#   Cloudflare One Client on headless Linux machines):
#     a. Zero Trust -> Access controls -> Service credentials
#        -> Service Tokens -> Create Service Token,记下 Client ID/Secret
#     b. Team & Resources -> Devices -> Management 标签页
#        -> Device enrollment permissions -> Manage -> Create new policy:
#           Action = Service Auth, Include: Service Token = <token 名>
#        保存后把该 policy 加入 enrollment permissions 并 Save
#   然后运行:
#     CF_TEAM=团队名 CF_ACCESS_CLIENT_ID=xxxx.access \
#     CF_ACCESS_CLIENT_SECRET=yyyy bash setup_warp_zerotrust.sh
#   (已持有注册 JWT 时可用 WARP_JWT=eyJ... 代替 service token 两个变量)
#
# 交互模式: 不带环境变量直接运行时,会依次询问团队名/凭据(从 /dev/tty
#   读取,兼容 wget|bash 管道方式);团队名直接回车即为免费版,凭据输入
#   不完整自动回退免费版。环境变量优先于交互输入,全部给齐则不询问。
#
# 可选环境变量:
#   ALLOWED_IPS=0.0.0.0/0   进隧道的网段。仪表盘的 split tunnel 设置对纯
#                           WireGuard 不生效,以这里为准;只做 Tunnel 私网
#                           访问可填私网段如 10.0.0.0/8
#   MTU=1280                NAT64 + WireGuard 双重开销下的安全值
#   ACTION=delete           反向操作: 注销本设备(释放团队设备席位)并卸载 wg0
#
# 关键实现依据(2026-08 核实, 参考 wgcf v2.2.32 / Aether / fscarmen api.sh):
#   - 注册接口返回平铺 JSON,没有 .result 包裹:
#       .id / .token / .config.interface.addresses.{v4,v6} / .config.peers[0]...
#   - 请求必须带 CF-Client-Version 头,否则被拒;可用组合:
#       v0a4471 + a-6.35-4471 (Aether), v0a1922 + a-6.3-1922 (wgcf)
#   - TLS 指纹需接近官方 App: 锁 TLS1.2 + HTTP/1.1,否则 error 1020
#   - 团队注册: 先用 service token 向 https://<team>.cloudflareaccess.com/warp
#     换 JWT(在 CF_Authorization cookie 或页面 token= 参数里),注册时带
#     CF-Access-Jwt-Assertion: <jwt> 头
#   - 403 = 本机 IP 被 Cloudflare 标记, 429 = 注册过于频繁,换 IP/稍后重试
#   - 内核无 wireguard 模块时(LXC/老内核)自动回退 wireguard-go 用户态实现
#     (wg-quick 内置该回退机制, 自启同样生效)
#
# 安全说明(供应链审计要点):
#   - 仅经 apt 官方源装包(wireguard-tools/curl/jq/wireguard-go), 不添加任何
#     第三方软件源, 不下载/执行任何远程二进制或脚本
#   - 外联仅 Cloudflare 官方域名; 连通性测试优先用官方 cdn-cgi/trace,
#     第三方 ip.sb 仅作兜底
#   - 安装阶段临时用公共 DNS64(nat64.net/Trex), 可用 DNS64_SERVERS 覆盖;
#     apt 有官方签名校验, API 全程 TLS, 结束后自动还原 DNS
#   - 私钥/凭据仅存本机 /etc/wireguard/(600), 只有公钥提交给注册接口
# ============================================================

# 0. 确保 root 权限
if [ "$EUID" -ne 0 ]; then
    echo "Error: Please run as root"
    exit 1
fi
umask 077

CF_TEAM="${CF_TEAM:-}"
CF_ACCESS_CLIENT_ID="${CF_ACCESS_CLIENT_ID:-}"
CF_ACCESS_CLIENT_SECRET="${CF_ACCESS_CLIENT_SECRET:-}"
WARP_JWT="${WARP_JWT:-}"
ALLOWED_IPS="${ALLOWED_IPS:-0.0.0.0/0}"
MTU="${MTU:-1280}"
ACTION="${ACTION:-setup}"
# 安装阶段临时使用的 DNS64 服务器(公共 NAT64 提供商无"官方"可言,可自行换成自建的)
DNS64_SERVERS="${DNS64_SERVERS:-2a01:4f8:c2c:123f::1 2001:67c:2b0::4}"

ACCOUNT_FILE=/etc/wireguard/warp-account.json
WG_CONF=/etc/wireguard/wg0.conf
API_HOST=https://api.cloudflareclient.com
# 依次尝试的 API 版本组合: "api版本|client版本"
API_VERSIONS=("v0a4471|a-6.35-4471" "v0a1922|a-6.3-1922")

TMPD=$(mktemp -d)

# ACTION=delete: 注销设备并卸载,不需要动 DNS
if [ "$ACTION" = "delete" ]; then
    if [ ! -s "$ACCOUNT_FILE" ]; then
        echo "No saved device ($ACCOUNT_FILE not found), nothing to delete."
    else
        command -v jq >/dev/null || { echo "Error: jq is required for delete"; exit 1; }
        DEL_ID=$(jq -r '.id' "$ACCOUNT_FILE")
        DEL_TOKEN=$(jq -r '.token' "$ACCOUNT_FILE")
        CODE=$(curl -sS -6 -m 20 -o /dev/null -w '%{http_code}' -X DELETE \
            -H 'User-Agent: okhttp/3.12.1' -H 'CF-Client-Version: a-6.35-4471' \
            -H "Authorization: Bearer $DEL_TOKEN" "$API_HOST/v0a4471/reg/$DEL_ID" || true)
        [ "$CODE" = "204" ] && echo "Device $DEL_ID deleted." || echo "Delete returned HTTP $CODE (device may already be gone)."
        rm -f "$ACCOUNT_FILE"
    fi
    systemctl disable --now wg-quick@wg0 2>/dev/null || true
    rm -f "$WG_CONF"
    echo "Uninstall done."
    exit 0
fi

# 交互式收集配置: 环境变量已提供的项不询问;无终端(非交互)时全部跳过
# 注意: 必须用"能否打开 /dev/tty"判断交互性, [ -r /dev/tty ] 在无控制终端时
# 也返回真,真正打开才报 ENXIO
interactive() { (: < /dev/tty) 2>/dev/null; }

ask() {
    # ask <变量名> <提示语> [1=密文输入]
    local var="$1" prompt="$2" secret="${3:-}"
    [ -n "${!var}" ] && return 0
    interactive || return 0
    if [ "$secret" = "1" ]; then
        read -rsp "$prompt" "$var" < /dev/tty || true
        echo > /dev/tty || true
    else
        read -rp "$prompt" "$var" < /dev/tty || true
    fi
}

if [ -z "$CF_TEAM" ] && interactive; then
    echo "------------------------------------------------------------"
    echo "直接回车 = 免费版 WARP;或输入 Zero Trust 团队名接入团队设备"
    echo "(团队模式需先在仪表盘创建 Service Token 并配置 Service Auth"
    echo " 注册策略,步骤见脚本头部注释)"
    echo "------------------------------------------------------------"
fi
ask CF_TEAM "Zero Trust 团队名(回车跳过): "
# 团队名归一化: 允许填 myteam / myteam.cloudflareaccess.com / https://myteam.cloudflareaccess.com/warp
CF_TEAM=$(echo "$CF_TEAM" | tr 'A-Z' 'a-z' | sed -E 's#^https?://##; s#\.cloudflareaccess\.com.*$##; s#/.*$##; s#[./]$##')
if [ -n "$CF_TEAM" ]; then
    ask CF_ACCESS_CLIENT_ID "Service Token Client ID (形如 xxxx.access): "
    ask CF_ACCESS_CLIENT_SECRET "Service Token Client Secret (输入不显示): " 1
    if [ -z "$WARP_JWT" ] && { [ -z "$CF_ACCESS_CLIENT_ID" ] || [ -z "$CF_ACCESS_CLIENT_SECRET" ]; }; then
        echo "*** 提示: 团队凭据不完整,本次按免费版 WARP 注册 ***"
        CF_TEAM=""
    fi
fi

# 1. 接管 DNS,并设置退出/报错自动恢复机制。按 resolv.conf 的实际形态分两种方式:
#   a) resolv.conf 是指向 /run/systemd/resolve/* 的符号链接(Ubuntu 默认 stub 模式):
#      文件内容对 resolved 没有作用,硬写只会踩坏 /run 里的 stub 文件。
#      改用 resolved.conf.d drop-in 注入全局 DNS64 + 路由域 ~.,恢复时删除并重启 resolved。
#      供应商/netplan 下发的原始 DNS 全程不动,也不存在快照过期问题。
#   b) 其余情况(foreign 普通文件 / NM / resolvconf / 无 resolved):
#      备份文件内容后改写,恢复时原位写回(cat > 保留 inode 与符号链接)。
#      禁止整文件 cp 替换: 换掉 inode 后 systemd-resolved 的 inotify 监控失效,
#      会出现 resolv.conf 内容正确但 resolvectl 看不到任何 DNS 的假象。
DNS_MODE="file"
RESOLV_BAK="$TMPD/resolv.conf.bak"
RESOLVED_DROPIN=/etc/systemd/resolved.conf.d/zz-warp-dns64.conf
RESOLV_REAL=$(readlink -f /etc/resolv.conf 2>/dev/null || echo /etc/resolv.conf)
if [ "${RESOLV_REAL#/run/systemd/resolve/}" != "$RESOLV_REAL" ] \
    && systemctl is-active --quiet systemd-resolved 2>/dev/null; then
    DNS_MODE=resolved
fi
set_dns64() {
    if [ "$DNS_MODE" = resolved ]; then
        mkdir -p /etc/systemd/resolved.conf.d
        printf '[Resolve]\nDNS=%s\nDomains=~.\n' "$DNS64_SERVERS" > "$RESOLVED_DROPIN"
        systemctl restart systemd-resolved
    else
        # -L 穿透符号链接备份内容(NM/resolvconf 管理时备份的是目标文件内容);
        # 备份失败(如悬空链接)按"原本不存在"处理,恢复时删除
        if [ -e /etc/resolv.conf ] || [ -L /etc/resolv.conf ]; then
            cp -L /etc/resolv.conf "$RESOLV_BAK" 2>/dev/null || : > "$RESOLV_BAK.missing"
        else
            : > "$RESOLV_BAK.missing"
        fi
        for ns in $DNS64_SERVERS; do echo "nameserver $ns"; done > /etc/resolv.conf
    fi
}
restore_dns() {
    echo "Restoring default DNS configuration..."
    if [ "$DNS_MODE" = resolved ]; then
        rm -f "$RESOLVED_DROPIN"
        systemctl restart systemd-resolved 2>/dev/null || true
    elif [ -f "$RESOLV_BAK.missing" ]; then
        rm -f /etc/resolv.conf "$RESOLV_BAK.missing"
    elif [ -f "$RESOLV_BAK" ]; then
        cat "$RESOLV_BAK" > /etc/resolv.conf 2>/dev/null || cp "$RESOLV_BAK" /etc/resolv.conf || true
    fi
}
cleanup() { restore_dns; rm -rf "$TMPD"; }
trap cleanup EXIT INT TERM

# 2. 临时启用 DNS64 保证初始化访问(apt 源大多只有 IPv4)
set_dns64

# 3. 安装轻量依赖(仅约 2MB)
APT_UPDATED=0
if ! command -v wg >/dev/null || ! command -v jq >/dev/null || ! command -v curl >/dev/null; then
    apt update -y && APT_UPDATED=1
    apt install -y wireguard-tools curl jq
fi

# 4. 团队模式: 用 service token 换注册 JWT
if [ -n "$CF_TEAM" ] && [ -z "$WARP_JWT" ]; then
    echo "Fetching enrollment JWT from https://${CF_TEAM}.cloudflareaccess.com/warp ..."
    JWT_CODE=$(curl -sS -6 -m 20 --max-redirs 0 -D "$TMPD/hdrs" -o "$TMPD/body" -w '%{http_code}' \
        -H "CF-Access-Client-Id: $CF_ACCESS_CLIENT_ID" \
        -H "CF-Access-Client-Secret: $CF_ACCESS_CLIENT_SECRET" \
        "https://${CF_TEAM}.cloudflareaccess.com/warp" || true)
    WARP_JWT=$(tr -d '\r' < "$TMPD/hdrs" | sed -n 's/^[Ss]et-[Cc]ookie:.*CF_Authorization=\([^;]*\).*/\1/p' | head -1)
    [ -z "$WARP_JWT" ] && WARP_JWT=$(grep -oE 'token=eyJ[A-Za-z0-9_.-]+' "$TMPD/body" | head -1 | cut -d= -f2-)
    if [ -z "$WARP_JWT" ]; then
        JWT_LOC=$(tr -d '\r' < "$TMPD/hdrs" | sed -n 's/^[Ll]ocation: *//p' | head -1)
        echo "Error: no enrollment JWT returned (HTTP ${JWT_CODE:-000}${JWT_LOC:+, Location: $JWT_LOC})"
        [ "$JWT_CODE" = "404" ] && echo "!! HTTP 404 = 团队域名不存在,团队名拼写有误(本次请求: ${CF_TEAM}.cloudflareaccess.com)"
        echo "排查清单:"
        echo "  1. 注册策略的 Action 必须是 Service Auth(选成 Allow 会出现"
        echo "     302 跳登录页,永远拿不到 JWT)"
        echo "  2. 策略建好后要回到 Device enrollment permissions 添加进列表并 Save"
        echo "  3. 检查 Client ID(xxx.access)/Secret 是否复制完整、token 是否过期"
        exit 1
    fi
    echo "JWT acquired."
fi

# 5. 生成 WireGuard 密钥
PRIVKEY=$(wg genkey)
PUBKEY=$(echo "$PRIVKEY" | wg pubkey)

# 5.5 探测内核 WireGuard 支持; 无模块时回退用户态 wireguard-go
# (wg-quick 内置该回退: ip link add 失败且装有 wireguard-go 时自动以用户态运行,
#  开机自启同样生效; LXC/自定义老内核 VPS 常见此场景)
USERSPACE_WG=0
if ! ip link add wgprobe type wireguard 2>/dev/null; then
    modprobe wireguard 2>/dev/null || true
fi
if ip link add wgprobe type wireguard 2>/dev/null; then
    ip link delete wgprobe 2>/dev/null || true
else
    echo "No kernel WireGuard module, falling back to userspace wireguard-go..."
    if [ ! -c /dev/net/tun ]; then
        echo "Error: /dev/net/tun not available (LXC 容器请先在面板/宿主机开启 TUN)"
        exit 1
    fi
    if ! command -v wireguard-go >/dev/null; then
        [ "$APT_UPDATED" = "0" ] && { apt update -y && APT_UPDATED=1; }
        apt install -y wireguard-go
    fi
    USERSPACE_WG=1
fi

RESP="$TMPD/reg.json"

# 6a. 已有设备凭据且模式一致时,用 PATCH 换新密钥(避免重复注册占用团队设备席位)
if [ -s "$ACCOUNT_FILE" ]; then
    OLD_ID=$(jq -r '.id' "$ACCOUNT_FILE")
    OLD_TOKEN=$(jq -r '.token' "$ACCOUNT_FILE")
    SAVED_TEAM=$(jq -r '.team // empty' "$ACCOUNT_FILE")
    if [ "$SAVED_TEAM" != "$CF_TEAM" ]; then
        echo "Mode changed (saved: '${SAVED_TEAM:-free}' -> current: '${CF_TEAM:-free}'), re-registering..."
        rm -f "$ACCOUNT_FILE"
    else
        echo "Existing device $OLD_ID found, updating its key instead of re-registering..."
        CODE=$(curl -sS -6 -m 25 --tlsv1.2 --tls-max 1.2 --http1.1 -X PATCH \
            -H 'User-Agent: okhttp/3.12.1' -H 'CF-Client-Version: a-6.35-4471' -H 'Content-Type: application/json' \
            -H "Authorization: Bearer $OLD_TOKEN" \
            -d "{\"key\":\"$PUBKEY\",\"key_type\":\"curve25519\",\"tunnel_type\":\"wireguard\"}" \
            -o "$RESP" -w '%{http_code}' "$API_HOST/v0a4471/reg/$OLD_ID" || true)
        if [ "$CODE" = "200" ] && jq -e '.config.interface.addresses.v4' "$RESP" >/dev/null 2>&1; then
            DEVICE_ID="$OLD_ID"
            DEVICE_TOKEN=$(jq -r '.token // empty' "$RESP")
            [ -z "$DEVICE_TOKEN" ] && DEVICE_TOKEN="$OLD_TOKEN"
            echo "Device key updated."
        else
            echo "Update failed (HTTP $CODE), falling back to fresh registration: $(head -c 200 "$RESP" 2>/dev/null)"
            rm -f "$ACCOUNT_FILE"
        fi
    fi
fi

# 6b. 全新注册
if [ ! -s "$ACCOUNT_FILE" ]; then
    if [ -n "$CF_TEAM" ]; then
        # 团队注册字段(参照 Aether TeamRegistration)
        INSTALL_ID=$(head -c 100 /dev/urandom | tr -dc 'A-Za-z0-9' | head -c 22)
        FCM_TOKEN="${INSTALL_ID}:APA91b$(head -c 300 /dev/urandom | tr -dc 'A-Za-z0-9' | head -c 134)"
        BODY=$(jq -nc --arg key "$PUBKEY" --arg tos "$(date -u +%FT%T.000Z)" --arg iid "$INSTALL_ID" \
            --arg fcm "$FCM_TOKEN" --arg name "$(hostname)" \
            '{key:$key, install_id:$iid, fcm_token:$fcm, tos:$tos, model:"PC", name:$name, serial_number:$iid, locale:"en_US"}')
    else
        # 免费版注册字段(参照 wgcf + Aether)
        SERIAL=$(head -c 8 /dev/urandom | od -An -tx1 | tr -d ' \n')
        BODY=$(jq -nc --arg key "$PUBKEY" --arg tos "$(date -u +%FT%T.000Z)" --arg sn "$SERIAL" \
            '{key:$key, install_id:"", fcm_token:"", tos:$tos, model:"PC", serial_number:$sn, os_version:"", key_type:"curve25519", tunnel_type:"wireguard", locale:"en_US", type:"Android"}')
    fi

    OK=""
    for combo in "${API_VERSIONS[@]}"; do
        VER=${combo%%|*}; CVER=${combo##*|}
        echo "Registering via $API_HOST/$VER/reg (CF-Client-Version: $CVER) ..."
        CODE=$(curl -sS -6 -m 25 --tlsv1.2 --tls-max 1.2 --http1.1 -X POST \
            -H 'User-Agent: okhttp/3.12.1' -H "CF-Client-Version: $CVER" -H 'Content-Type: application/json' \
            ${WARP_JWT:+-H "CF-Access-Jwt-Assertion: $WARP_JWT"} \
            -d "$BODY" -o "$RESP" -w '%{http_code}' "$API_HOST/$VER/reg" || true)
        if [ "$CODE" = "200" ] && jq -e '.id and .config.interface.addresses.v4' "$RESP" >/dev/null 2>&1; then
            OK=1
            break
        fi
        echo "  failed, HTTP ${CODE:-000}: $(head -c 300 "$RESP" 2>/dev/null)"
    done
    if [ -z "$OK" ]; then
        echo "Error: Cloudflare WARP registration failed on all API versions."
        echo "Hint: 403 = VPS IP flagged by Cloudflare; 429 = rate limited. Try later or from another IP."
        exit 1
    fi
    DEVICE_ID=$(jq -r '.id' "$RESP")
    DEVICE_TOKEN=$(jq -r '.token' "$RESP")
    jq -nc --arg id "$DEVICE_ID" --arg token "$DEVICE_TOKEN" --arg team "$CF_TEAM" \
        '{id:$id, token:$token, team:$team}' > "$ACCOUNT_FILE"
    echo "Registered device $DEVICE_ID."
fi

# 7. 解析配置(注意: 返回是平铺 JSON,没有 .result 包裹)
CF_IPV4=$(jq -r '.config.interface.addresses.v4 // empty' "$RESP")
CF_IPV6=$(jq -r '.config.interface.addresses.v6 // empty' "$RESP")
PEER_KEY=$(jq -r '.config.peers[0].public_key // empty' "$RESP")
[ -z "$PEER_KEY" ] && PEER_KEY="bmXOC+F1FxEMF9dyiK2H5/1SUtzH0JuVo51h2wPfgyo="
EP6=$(jq -r '.config.peers[0].endpoint.v6 // empty' "$RESP")
# endpoint.v6 可能是 [addr]:port 或裸地址, 统一剥成裸地址
case "$EP6" in
    \[*\]*) EP6="${EP6#[}"; EP6="${EP6%%\]*}" ;;
esac
[ -z "$EP6" ] && EP6="2606:4700:d0::a29f:c001"  # engage.cloudflareclient.com 的 IPv6
ORG=$(jq -r '.account.organization // empty' "$RESP")

# 8. 写入 wg0.conf(默认纯 IPv4 出口,不碰 IPv6 路由)
mkdir -p /etc/wireguard
cat << EOF > "$WG_CONF"
[Interface]
PrivateKey = $PRIVKEY
Address = ${CF_IPV4}/32${CF_IPV6:+, ${CF_IPV6}/128}
MTU = $MTU

[Peer]
PublicKey = $PEER_KEY
AllowedIPs = $ALLOWED_IPS
Endpoint = [$EP6]:2408
PersistentKeepalive = 25
EOF
chmod 600 "$WG_CONF"

# 9. 启动接口并开启开机自启
systemctl enable wg-quick@wg0
systemctl restart wg-quick@wg0

# 触发 trap 还原 VPS 默认 DNS 后进行网络验证
restore_dns
trap - EXIT INT TERM
rm -rf "$TMPD"

echo "--------------------------------------------------"
echo "Setup complete!${ORG:+ Organization: $ORG}${CF_TEAM:+ (team: $CF_TEAM)}"
echo "Testing connectivity with default DNS:"
# 优先用 Cloudflare 官方 trace 取出口 IP(第一方),第三方 ip.sb 仅作兜底
V4EXIT=$(curl -4 -s --connect-timeout 5 https://www.cloudflare.com/cdn-cgi/trace | sed -n 's/^ip=//p')
[ -z "$V4EXIT" ] && V4EXIT=$(curl -4 -s --connect-timeout 5 ip.sb 2>/dev/null || echo "Failed")
V6EXIT=$(curl -6 -s --connect-timeout 5 https://www.cloudflare.com/cdn-cgi/trace | sed -n 's/^ip=//p')
[ -z "$V6EXIT" ] && V6EXIT=$(curl -6 -s --connect-timeout 5 ip.sb 2>/dev/null || echo "Failed")
echo "IPv4 Exit: $V4EXIT"
echo "IPv6 Exit: $V6EXIT"
echo "WARP status: $(curl -4 -s --connect-timeout 5 https://www.cloudflare.com/cdn-cgi/trace | grep '^warp=' || echo 'unknown')"
[ "$USERSPACE_WG" = "1" ] && echo "Note: 当前为用户态 wireguard-go(无内核模块),速度略慢;如需内核级性能请换 >=5.6 内核的系统镜像"
echo "--------------------------------------------------"