#!/bin/bash

# ============================================================
#  TikTok 流量分流脚本
#  在台湾 VPS 上运行，让 TikTok 流量走美国 VPS
#  原理: SSH 动态隧道(SOCKS5) + redsocks + iptables 分流
#  适用: Debian / Ubuntu  |  认证方式: 密码
# ============================================================

set -e

RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
CYAN="\033[36m"
BOLD="\033[1m"
NC="\033[0m"

# ============================================================
# ⚙️  请修改以下配置
# ============================================================
US_VPS_IP="142.171.202.195"   # 美国 VPS IP
US_VPS_PORT="22"               # 美国 VPS SSH 端口
US_VPS_USER="root"             # 美国 VPS 用户名
US_VPS_PASS="2miijAPMZB"    # 美国 VPS 密码 ← 改这里
SOCKS5_PORT="7891"             # 本地 SOCKS5 监听端口
# ============================================================

TIKTOK_DOMAINS=(
    "tiktok.com"
    "www.tiktok.com"
    "m.tiktok.com"
    "live.tiktok.com"
    "api.tiktok.com"
    "tiktokv.com"
    "ttwstatic.com"
    "tiktokcdn.com"
    "tiktokcdn-us.com"
    "bytedance.com"
    "ibytedtos.com"
    "muscdn.com"
    "musical.ly"
)

IPSET_NAME="tiktok_ips"
REDSOCKS_CONF="/etc/redsocks.conf"

log_info()  { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn()  { echo -e "${YELLOW}[!]${NC} $1"; }
log_error() { echo -e "${RED}[✗]${NC} $1"; exit 1; }
log_title() { echo -e "\n${CYAN}${BOLD}─── $1 ───────────────────────────────────${NC}"; }

check_root() {
    [ "$EUID" -ne 0 ] && log_error "请使用 root 权限运行此脚本"
}

install_deps() {
    log_title "安装依赖"
    apt-get update -qq
    apt-get install -y redsocks ipset iptables dnsutils curl sshpass -qq
    log_info "依赖安装完成（含 sshpass）"
}

test_ssh() {
    log_title "测试 SSH 连接"
    if sshpass -p "$US_VPS_PASS" ssh \
        -o StrictHostKeyChecking=no \
        -o ConnectTimeout=10 \
        -p "$US_VPS_PORT" \
        "${US_VPS_USER}@${US_VPS_IP}" "echo ok" &>/dev/null; then
        log_info "SSH 连接美国 VPS 成功"
    else
        log_error "SSH 连接失败，请检查 IP / 用户名 / 密码是否正确"
    fi
}

setup_redsocks() {
    log_title "配置 redsocks"
    cat > "$REDSOCKS_CONF" << EOF
base {
    log_debug = off;
    log_info = on;
    log = "syslog:daemon";
    daemon = on;
    redirector = iptables;
}

redsocks {
    local_ip = 127.0.0.1;
    local_port = 12345;
    ip = 127.0.0.1;
    port = ${SOCKS5_PORT};
    type = socks5;
}
EOF
    log_info "redsocks 配置完成"
}

create_tunnel_service() {
    log_title "创建 SSH 隧道 systemd 服务"

    # 把密码存到只有 root 可读的文件，避免 ps 命令暴露
    echo "$US_VPS_PASS" > /etc/tiktok_tunnel.pass
    chmod 600 /etc/tiktok_tunnel.pass

    cat > /etc/systemd/system/tiktok-tunnel.service << EOF
[Unit]
Description=TikTok SSH Tunnel to US VPS
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=root
Environment=SSHPASS_FILE=/etc/tiktok_tunnel.pass
ExecStart=/bin/bash -c 'sshpass -f /etc/tiktok_tunnel.pass ssh \
    -N \
    -D 127.0.0.1:${SOCKS5_PORT} \
    -o StrictHostKeyChecking=no \
    -o ServerAliveInterval=30 \
    -o ServerAliveCountMax=3 \
    -o ExitOnForwardFailure=yes \
    -o ConnectTimeout=15 \
    -p ${US_VPS_PORT} \
    ${US_VPS_USER}@${US_VPS_IP}'
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable tiktok-tunnel
    systemctl restart tiktok-tunnel
    sleep 3

    if systemctl is-active --quiet tiktok-tunnel; then
        log_info "SSH 隧道服务启动成功（SOCKS5 127.0.0.1:${SOCKS5_PORT}）"
    else
        log_error "SSH 隧道启动失败，请运行: journalctl -u tiktok-tunnel -n 30"
    fi
}

resolve_tiktok_ips() {
    log_title "解析 TikTok IP 并写入 ipset"
    ipset destroy "$IPSET_NAME" 2>/dev/null || true
    ipset create "$IPSET_NAME" hash:ip timeout 86400

    local total=0
    for domain in "${TIKTOK_DOMAINS[@]}"; do
        local ips
        ips=$(dig +short "$domain" A 2>/dev/null | grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$')
        for ip in $ips; do
            ipset add "$IPSET_NAME" "$ip" 2>/dev/null || true
            ((total++)) || true
        done
    done
    log_info "共解析 ${total} 条 TikTok IP"
}

setup_iptables() {
    log_title "配置 iptables 透明代理规则"

    # 清理旧规则
    iptables -t nat -D OUTPUT    -m set --match-set "$IPSET_NAME" dst -p tcp -j REDIRECT --to-port 12345 2>/dev/null || true
    iptables -t nat -D PREROUTING -m set --match-set "$IPSET_NAME" dst -p tcp -j REDIRECT --to-port 12345 2>/dev/null || true

    # 本机发出的流量
    iptables -t nat -A OUTPUT \
        -m set --match-set "$IPSET_NAME" dst \
        -p tcp -j REDIRECT --to-port 12345

    # 经过本机转发的流量（如作为网关）
    iptables -t nat -A PREROUTING \
        -m set --match-set "$IPSET_NAME" dst \
        -p tcp -j REDIRECT --to-port 12345

    log_info "iptables 规则设置完成"
}

start_redsocks() {
    log_title "启动 redsocks"
    pkill redsocks 2>/dev/null || true
    sleep 1
    redsocks -c "$REDSOCKS_CONF"
    sleep 1
    if pgrep -x redsocks > /dev/null; then
        log_info "redsocks 启动成功（监听 127.0.0.1:12345）"
    else
        log_error "redsocks 启动失败"
    fi
}

save_and_persist() {
    log_title "保存规则并设置开机自启"

    # 保存 ipset
    ipset save > /etc/ipset.conf

    # 保存 iptables
    apt-get install -y iptables-persistent -qq 2>/dev/null || true
    iptables-save > /etc/iptables/rules.v4 2>/dev/null || true

    # ipset 开机恢复
    cat > /etc/systemd/system/ipset-restore.service << EOF
[Unit]
Description=Restore ipset on boot
Before=network.target

[Service]
Type=oneshot
ExecStart=/usr/sbin/ipset restore -f /etc/ipset.conf
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
    systemctl daemon-reload
    systemctl enable ipset-restore

    # redsocks 开机启动
    cat > /etc/systemd/system/redsocks-tiktok.service << EOF
[Unit]
Description=Redsocks for TikTok
After=tiktok-tunnel.service ipset-restore.service

[Service]
Type=forking
ExecStart=/usr/sbin/redsocks -c /etc/redsocks.conf
Restart=always

[Install]
WantedBy=multi-user.target
EOF
    systemctl daemon-reload
    systemctl enable redsocks-tiktok

    # 每6小时刷新 TikTok IP
    cat > /usr/local/bin/tiktok_refresh_ips.sh << EOF
#!/bin/bash
DOMAINS="${TIKTOK_DOMAINS[*]}"
for domain in \$DOMAINS; do
    ips=\$(dig +short "\$domain" A 2>/dev/null | grep -E '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+\$')
    for ip in \$ips; do
        ipset add "$IPSET_NAME" "\$ip" 2>/dev/null || true
    done
done
ipset save > /etc/ipset.conf
EOF
    chmod +x /usr/local/bin/tiktok_refresh_ips.sh
    echo "0 */6 * * * root /usr/local/bin/tiktok_refresh_ips.sh" > /etc/cron.d/tiktok-refresh

    log_info "开机自启配置完成，每6小时自动刷新 TikTok IP"
}

verify() {
    log_title "验证 TikTok 访问"
    echo -n "  正在通过美国 VPS 测试..."

    local result
    result=$(curl --socks5 127.0.0.1:"$SOCKS5_PORT" \
        -s --max-time 15 \
        -A "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36" \
        -H "Accept-Encoding: gzip" \
        "https://www.tiktok.com" 2>/dev/null | gunzip 2>/dev/null)

    local region
    region=$(echo "$result" | grep -o '"region":"[^"]*"' | head -1 | cut -f4 -d'"')

    echo ""
    if [ -n "$region" ]; then
        log_info "TikTok 访问成功！Region: ${GREEN}【${region}】${NC}"
    else
        log_warn "隧道已建立，但未获取到 region（TikTok 可能识别为 IDC IP，属正常现象）"
    fi
}

print_summary() {
    echo ""
    echo -e "${CYAN}${BOLD}╔══════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}${BOLD}║              🎉 配置完成！                      ║${NC}"
    echo -e "${CYAN}${BOLD}╚══════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "  ${GREEN}隧道路由:${NC}  台湾 VPS → SSH SOCKS5 → 美国 VPS (${US_VPS_IP})"
    echo -e "  ${GREEN}分流策略:${NC}  仅 TikTok 相关 IP 走美国，其他流量不变"
    echo ""
    echo -e "  ${YELLOW}常用命令:${NC}"
    echo -e "    查看隧道:    systemctl status tiktok-tunnel"
    echo -e "    重启隧道:    systemctl restart tiktok-tunnel"
    echo -e "    查看日志:    journalctl -u tiktok-tunnel -f"
    echo -e "    刷新IP列表:  bash /usr/local/bin/tiktok_refresh_ips.sh"
    echo -e "    查看IP列表:  ipset list $IPSET_NAME"
    echo ""
    echo -e "  ${YELLOW}卸载方法:${NC}"
    echo -e "    systemctl stop tiktok-tunnel redsocks-tiktok"
    echo -e "    systemctl disable tiktok-tunnel redsocks-tiktok ipset-restore"
    echo -e "    iptables -t nat -F OUTPUT && iptables -t nat -F PREROUTING"
    echo -e "    ipset destroy $IPSET_NAME"
    echo ""
}

# ── 主流程 ──────────────────────────────────────
check_root
install_deps
test_ssh
setup_redsocks
create_tunnel_service
resolve_tiktok_ips
setup_iptables
start_redsocks
save_and_persist
verify
print_summary