#!/bin/bash
# SS-Rust 一键安装脚本
# 支持: ss2022 (2022-blake3-aes-128-gcm) / ss128 (aes-128-gcm) / 双节点
# 自动生成 SS订阅 + Surge + Clash 配置
# GitHub: https://github.com/mango082888-bit/ss-rust

# ============ 颜色 ============
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

# ============ 基础检测 ============
check_root() {
    [[ $EUID -ne 0 ]] && error "请使用 root 用户运行"
}

get_pkg_manager() {
    if command -v apt &>/dev/null; then
        PKG="apt"
    elif command -v yum &>/dev/null; then
        PKG="yum"
    elif command -v apk &>/dev/null; then
        PKG="apk"
    else
        error "不支持的包管理器"
    fi
}

install_deps() {
    info "安装依赖..."
    case $PKG in
        apt) apt update -qq &>/dev/null; apt install -y -qq curl openssl xz-utils tar chrony &>/dev/null ;;
        yum) yum install -y -q curl openssl xz tar chrony &>/dev/null ;;
        apk) apk add --quiet curl openssl xz tar chrony &>/dev/null ;;
    esac
}

get_arch() {
    case $(uname -m) in
        x86_64)  echo "x86_64-unknown-linux-gnu" ;;
        aarch64) echo "aarch64-unknown-linux-gnu" ;;
        armv7l)  echo "armv7-unknown-linux-gnueabihf" ;;
        *)       error "不支持的架构: $(uname -m)" ;;
    esac
}

get_ip() {
    IP=$(curl -s4m5 ip.sb 2>/dev/null || curl -s4m5 ifconfig.me 2>/dev/null || curl -s4m5 ipinfo.io/ip 2>/dev/null)
    [[ -z "$IP" ]] && error "无法获取公网IP"
    echo "$IP"
}

# ============ 时间同步 ============
sync_time() {
    info "同步系统时间..."
    if command -v timedatectl &>/dev/null; then
        timedatectl set-ntp true 2>/dev/null || true
    fi
    if command -v chronyd &>/dev/null; then
        systemctl enable --now chronyd 2>/dev/null || true
        chronyc makestep 2>/dev/null || true
    fi
    info "当前时间: $(date '+%Y-%m-%d %H:%M:%S %Z')"
}

# ============ 安装 ss-rust ============
install_ssrust() {
    info "安装 shadowsocks-rust..."
    ARCH_NAME=$(get_arch)
    
    LATEST=$(curl -sLm10 https://api.github.com/repos/shadowsocks/shadowsocks-rust/releases/latest | grep tag_name | head -1 | grep -oP 'v[\d.]+')
    [[ -z "$LATEST" ]] && LATEST="v1.24.0"
    info "版本: $LATEST"
    
    URL="https://github.com/shadowsocks/shadowsocks-rust/releases/download/${LATEST}/shadowsocks-${LATEST}.${ARCH_NAME}.tar.xz"
    
    cd /tmp
    rm -f ss-rust-dl.tar.xz
    curl -sL "$URL" -o ss-rust-dl.tar.xz
    
    # 验证下载
    FILE_SIZE=$(stat -c%s ss-rust-dl.tar.xz 2>/dev/null || stat -f%z ss-rust-dl.tar.xz 2>/dev/null)
    [[ "$FILE_SIZE" -lt 100000 ]] && error "下载失败，文件太小 (${FILE_SIZE} bytes)，可能是网络问题"
    
    file ss-rust-dl.tar.xz | grep -q "XZ" || error "下载的文件不是有效的 XZ 压缩包"
    
    tar xf ss-rust-dl.tar.xz || error "解压失败"
    
    [[ ! -f ssserver ]] && error "解压后找不到 ssserver 二进制文件"
    
    cp -f ssserver /usr/local/bin/
    cp -f sslocal /usr/local/bin/ 2>/dev/null || true
    cp -f ssurl /usr/local/bin/ 2>/dev/null || true
    chmod +x /usr/local/bin/ss*
    
    # 清理
    rm -f ss-rust-dl.tar.xz ssserver sslocal ssurl ssmanager ssservice
    
    /usr/local/bin/ssserver --version && info "shadowsocks-rust 安装完成" || error "安装验证失败"
}

# ============ 节点选择菜单 ============
select_node_type() {
    echo ""
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo -e "${CYAN}  🔐 选择节点类型${NC}"
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo ""
    echo -e "  ${GREEN}1.${NC} SS2022-128  (2022-blake3-aes-128-gcm) — 新协议，推荐"
    echo -e "  ${GREEN}2.${NC} SS-AES-128  (aes-128-gcm) — 传统协议，兼容性好"
    echo -e "  ${GREEN}3.${NC} 双节点全装  (SS2022 + SS128)"
    echo ""
    read -p "请选择 [1-3] (默认3): " node_choice
    node_choice=${node_choice:-3}
}

# ============ 生成配置 ============
gen_config() {
    SERVER_IP=$(get_ip)
    mkdir -p /etc/shadowsocks-rust
    
    PORT_2022=$((RANDOM % 10000 + 20000))
    PORT_RAW=$((RANDOM % 10000 + 30000))
    KEY_2022=$(openssl rand -base64 16)
    KEY_RAW=$(openssl rand -base64 16)
    METHOD_2022="2022-blake3-aes-128-gcm"
    METHOD_RAW="aes-128-gcm"
    
    case "${node_choice}" in
        1)
            cat > /etc/shadowsocks-rust/config.json << EOF
{
    "servers": [
        {
            "server": "0.0.0.0",
            "server_port": ${PORT_2022},
            "method": "${METHOD_2022}",
            "password": "${KEY_2022}",
            "timeout": 300,
            "fast_open": true
        }
    ]
}
EOF
            NODE_MODE="ss2022"
            ;;
        2)
            cat > /etc/shadowsocks-rust/config.json << EOF
{
    "servers": [
        {
            "server": "0.0.0.0",
            "server_port": ${PORT_RAW},
            "method": "${METHOD_RAW}",
            "password": "${KEY_RAW}",
            "timeout": 300,
            "fast_open": true
        }
    ]
}
EOF
            NODE_MODE="ss128"
            ;;
        3|*)
            cat > /etc/shadowsocks-rust/config.json << EOF
{
    "servers": [
        {
            "server": "0.0.0.0",
            "server_port": ${PORT_2022},
            "method": "${METHOD_2022}",
            "password": "${KEY_2022}",
            "timeout": 300,
            "fast_open": true
        },
        {
            "server": "0.0.0.0",
            "server_port": ${PORT_RAW},
            "method": "${METHOD_RAW}",
            "password": "${KEY_RAW}",
            "timeout": 300,
            "fast_open": true
        }
    ]
}
EOF
            NODE_MODE="both"
            ;;
    esac
    
    info "配置文件: /etc/shadowsocks-rust/config.json"
}

# ============ systemd 服务 ============
setup_service() {
    cat > /etc/systemd/system/ss-rust.service << 'EOF'
[Unit]
Description=Shadowsocks-Rust Server
After=network.target

[Service]
Type=simple
ExecStart=/usr/local/bin/ssserver -c /etc/shadowsocks-rust/config.json
Restart=on-failure
RestartSec=5
LimitNOFILE=65535

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable --now ss-rust
    sleep 2
    
    if systemctl is-active --quiet ss-rust; then
        info "ss-rust 服务启动成功"
    else
        echo ""
        journalctl -u ss-rust -n 5 --no-pager
        error "ss-rust 服务启动失败"
    fi
}

# ============ 读取现有配置 ============
load_config() {
    if [[ ! -f /etc/shadowsocks-rust/config.json ]]; then
        return 1
    fi
    SERVER_IP=$(get_ip)
    
    # 用 python3 或 grep 解析
    if command -v python3 &>/dev/null; then
        eval $(python3 -c "
import json
with open('/etc/shadowsocks-rust/config.json') as f:
    c = json.load(f)
for i,s in enumerate(c['servers']):
    m = s['method']
    if '2022' in m:
        print(f'PORT_2022={s[\"server_port\"]}')
        print(f'KEY_2022={s[\"password\"]}')
        print(f'METHOD_2022={m}')
    else:
        print(f'PORT_RAW={s[\"server_port\"]}')
        print(f'KEY_RAW={s[\"password\"]}')
        print(f'METHOD_RAW={m}')
n = len(c['servers'])
if n == 1:
    if '2022' in c['servers'][0]['method']:
        print('NODE_MODE=ss2022')
    else:
        print('NODE_MODE=ss128')
else:
    print('NODE_MODE=both')
" 2>/dev/null)
    fi
    return 0
}

# ============ 生成订阅 ============
gen_subscribe() {
    load_config || return
    SUB_DIR="/etc/shadowsocks-rust/subscribe"
    mkdir -p "$SUB_DIR"
    
    SS_URIS=""
    SURGE_PROXIES=""
    CLASH_PROXIES=""
    INFO_NODES=""
    SURGE_NAMES=""
    
    if [[ -n "${PORT_2022:-}" ]]; then
        URI_2022="ss://$(echo -n "${METHOD_2022}:${KEY_2022}" | base64 -w0)@${SERVER_IP}:${PORT_2022}#SS2022-128"
        SS_URIS="${SS_URIS}${URI_2022}\n"
        SURGE_PROXIES="${SURGE_PROXIES}SS2022-128 = ss, ${SERVER_IP}, ${PORT_2022}, encrypt-method=${METHOD_2022}, password=${KEY_2022}\n"
        SURGE_NAMES="${SURGE_NAMES}SS2022-128, "
        CLASH_PROXIES="${CLASH_PROXIES}  - name: SS2022-128\n    type: ss\n    server: ${SERVER_IP}\n    port: ${PORT_2022}\n    cipher: ${METHOD_2022}\n    password: \"${KEY_2022}\"\n\n"
        INFO_NODES="${INFO_NODES}【SS2022-AES-128】新协议\n  地址: ${SERVER_IP}\n  端口: ${PORT_2022}\n  加密: ${METHOD_2022}\n  密码: ${KEY_2022}\n\n"
    fi
    
    if [[ -n "${PORT_RAW:-}" ]]; then
        URI_RAW="ss://$(echo -n "${METHOD_RAW}:${KEY_RAW}" | base64 -w0)@${SERVER_IP}:${PORT_RAW}#SS-AES-128"
        SS_URIS="${SS_URIS}${URI_RAW}\n"
        SURGE_PROXIES="${SURGE_PROXIES}SS-AES-128 = ss, ${SERVER_IP}, ${PORT_RAW}, encrypt-method=${METHOD_RAW}, password=${KEY_RAW}\n"
        SURGE_NAMES="${SURGE_NAMES}SS-AES-128, "
        CLASH_PROXIES="${CLASH_PROXIES}  - name: SS-AES-128\n    type: ss\n    server: ${SERVER_IP}\n    port: ${PORT_RAW}\n    cipher: ${METHOD_RAW}\n    password: \"${KEY_RAW}\"\n\n"
        INFO_NODES="${INFO_NODES}【SS-AES-128】传统协议\n  地址: ${SERVER_IP}\n  端口: ${PORT_RAW}\n  加密: ${METHOD_RAW}\n  密码: ${KEY_RAW}\n\n"
    fi
    
    # Base64 订阅
    echo -e "${SS_URIS}" | base64 -w0 > "$SUB_DIR/subscribe.txt"
    
    # Surge
    cat > "$SUB_DIR/surge.conf" << EOF
# Surge SS 配置 | $(date '+%Y-%m-%d %H:%M:%S') | ${SERVER_IP}
[Proxy]
$(echo -e "${SURGE_PROXIES}")
[Proxy Group]
Proxy = select, $(echo "${SURGE_NAMES}DIRECT" | sed 's/, DIRECT$/, DIRECT/')
EOF
    
    # Clash
    cat > "$SUB_DIR/clash.yaml" << EOF
# Clash SS 配置 | $(date '+%Y-%m-%d %H:%M:%S')
proxies:
$(echo -e "${CLASH_PROXIES}")
EOF
    
    # Info
    cat > "$SUB_DIR/info.txt" << EOF
════════════════════════════════════════
  Shadowsocks-Rust 节点信息
  $(date '+%Y-%m-%d %H:%M:%S') | ${SERVER_IP}
════════════════════════════════════════

$(echo -e "${INFO_NODES}")
【SS 链接】
$(echo -e "${SS_URIS}")
════════════════════════════════════════
EOF
}

# ============ 显示结果 ============
show_result() {
    load_config || return
    echo ""
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo -e "${CYAN}  🚀 Shadowsocks-Rust 安装完成${NC}"
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo ""
    
    if [[ -n "${PORT_2022:-}" ]]; then
        echo -e "${GREEN}【节点1】SS2022-AES-128 (新协议)${NC}"
        echo -e "  地址: ${SERVER_IP}"
        echo -e "  端口: ${YELLOW}${PORT_2022}${NC}"
        echo -e "  加密: ${METHOD_2022}"
        echo -e "  密码: ${YELLOW}${KEY_2022}${NC}"
        echo ""
    fi
    
    if [[ -n "${PORT_RAW:-}" ]]; then
        echo -e "${GREEN}【节点2】SS-AES-128 (传统协议)${NC}"
        echo -e "  地址: ${SERVER_IP}"
        echo -e "  端口: ${YELLOW}${PORT_RAW}${NC}"
        echo -e "  加密: ${METHOD_RAW}"
        echo -e "  密码: ${YELLOW}${KEY_RAW}${NC}"
        echo ""
    fi
    
    echo -e "${CYAN}────────────────────────────────────────${NC}"
    echo -e "${GREEN}【Surge 格式】${NC}"
    [[ -n "${PORT_2022:-}" ]] && echo -e "  SS2022-128 = ss, ${SERVER_IP}, ${PORT_2022}, encrypt-method=${METHOD_2022}, password=${KEY_2022}"
    [[ -n "${PORT_RAW:-}" ]] && echo -e "  SS-AES-128 = ss, ${SERVER_IP}, ${PORT_RAW}, encrypt-method=${METHOD_RAW}, password=${KEY_RAW}"
    echo ""
    
    echo -e "${CYAN}────────────────────────────────────────${NC}"
    echo -e "${GREEN}【SS 链接】${NC}"
    [[ -n "${URI_2022:-}" ]] && echo -e "  ${URI_2022}"
    [[ -n "${URI_RAW:-}" ]] && echo -e "  ${URI_RAW}"
    echo ""
    
    echo -e "${CYAN}────────────────────────────────────────${NC}"
    echo -e "${GREEN}【文件位置】${NC}"
    echo -e "  配置: /etc/shadowsocks-rust/config.json"
    echo -e "  订阅: /etc/shadowsocks-rust/subscribe/subscribe.txt"
    echo -e "  Surge: /etc/shadowsocks-rust/subscribe/surge.conf"
    echo -e "  Clash: /etc/shadowsocks-rust/subscribe/clash.yaml"
    echo ""
    echo -e "${CYAN}【管理】${NC} 再次运行本脚本即可进入管理菜单"
    echo -e "${CYAN}════════════════════════════════════════${NC}"
}

# ============ 查看配置 ============
show_config() {
    load_config || error "未安装 shadowsocks-rust"
    echo ""
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo -e "${CYAN}  📋 当前配置${NC}"
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo ""
    echo -e "${GREEN}服务状态:${NC} $(systemctl is-active ss-rust 2>/dev/null || echo '未运行')"
    echo ""
    cat /etc/shadowsocks-rust/config.json
    echo ""
    [[ -f /etc/shadowsocks-rust/subscribe/info.txt ]] && cat /etc/shadowsocks-rust/subscribe/info.txt
}

# ============ 修改端口 ============
change_port() {
    load_config || error "未安装 shadowsocks-rust"
    echo ""
    echo -e "${GREEN}当前节点:${NC}"
    [[ -n "${PORT_2022:-}" ]] && echo -e "  1) SS2022-128 | 端口 ${PORT_2022}"
    [[ -n "${PORT_RAW:-}" ]] && echo -e "  2) SS-AES-128 | 端口 ${PORT_RAW}"
    echo ""
    read -p "输入节点编号: " pn
    read -p "输入新端口: " new_port
    
    python3 -c "
import json
with open('/etc/shadowsocks-rust/config.json') as f:
    c = json.load(f)
idx = int('${pn}') - 1
if 0 <= idx < len(c['servers']):
    c['servers'][idx]['server_port'] = int('${new_port}')
    with open('/etc/shadowsocks-rust/config.json','w') as f:
        json.dump(c, f, indent=4)
" 2>/dev/null
    
    systemctl restart ss-rust
    gen_subscribe
    info "端口已修改为 ${new_port}，订阅已更新"
}

# ============ 重置密钥 ============
reset_keys() {
    load_config || error "未安装 shadowsocks-rust"
    
    python3 -c "
import json, subprocess, base64, os
with open('/etc/shadowsocks-rust/config.json') as f:
    c = json.load(f)
for s in c['servers']:
    s['password'] = base64.b64encode(os.urandom(16)).decode()
with open('/etc/shadowsocks-rust/config.json','w') as f:
    json.dump(c, f, indent=4)
" 2>/dev/null
    
    systemctl restart ss-rust
    gen_subscribe
    info "密钥已重置"
    show_result
}

# ============ 卸载 ============
uninstall() {
    warn "卸载 shadowsocks-rust..."
    systemctl stop ss-rust 2>/dev/null
    systemctl disable ss-rust 2>/dev/null
    rm -f /etc/systemd/system/ss-rust.service
    rm -f /usr/local/bin/ssserver /usr/local/bin/sslocal /usr/local/bin/ssurl
    rm -rf /etc/shadowsocks-rust
    systemctl daemon-reload
    info "卸载完成"
}

# ============ 完整安装流程 ============
do_install() {
    get_pkg_manager
    install_deps
    sync_time
    select_node_type
    install_ssrust
    gen_config
    setup_service
    gen_subscribe
    show_result
}

# ============ 管理菜单 ============
show_menu() {
    echo ""
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo -e "${CYAN}  🚀 SS-Rust 管理面板${NC}"
    echo -e "${CYAN}════════════════════════════════════════${NC}"
    echo ""
    echo -e "  ${GREEN}1.${NC} 重新安装 (选择节点类型)"
    echo -e "  ${GREEN}2.${NC} 查看配置 + 节点信息"
    echo -e "  ${GREEN}3.${NC} 修改端口"
    echo -e "  ${GREEN}4.${NC} 重置密钥"
    echo -e "  ${GREEN}5.${NC} 启动服务"
    echo -e "  ${GREEN}6.${NC} 停止服务"
    echo -e "  ${GREEN}7.${NC} 重启服务"
    echo -e "  ${GREEN}8.${NC} 查看日志"
    echo -e "  ${RED}9.${NC} 卸载"
    echo -e "  ${YELLOW}0.${NC} 退出"
    echo ""
    read -p "请选择 [0-9]: " choice
    
    case "$choice" in
        1) do_install ;;
        2) show_config ;;
        3) change_port ;;
        4) reset_keys ;;
        5) systemctl start ss-rust && info "服务已启动" ;;
        6) systemctl stop ss-rust && info "服务已停止" ;;
        7) systemctl restart ss-rust && info "服务已重启" ;;
        8) journalctl -u ss-rust --no-pager -n 30 ;;
        9) uninstall ;;
        0) exit 0 ;;
        *) warn "无效选择" ;;
    esac
}

# ============ 主入口 ============
main() {
    check_root
    case "${1:-}" in
        install)  do_install ;;
        uninstall|remove) uninstall ;;
        show|config|info) show_config ;;
        restart)  systemctl restart ss-rust && info "服务已重启" ;;
        start)    systemctl start ss-rust && info "服务已启动" ;;
        stop)     systemctl stop ss-rust && info "服务已停止" ;;
        log|logs) journalctl -u ss-rust --no-pager -n 30 ;;
        reset)    reset_keys ;;
        *)
            if [[ -f /etc/shadowsocks-rust/config.json ]]; then
                show_menu
            else
                do_install
            fi
            ;;
    esac
}

main "$@"