#!/bin/bash

# 脚本: setup-cloudflare-ssh-access.sh
# 描述: 配置Cloudflare Tunnels和短期证书功能以启用Web SSH访问
# 使用方法: ./setup-cloudflare-ssh-access.sh

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 无颜色

# 变量声明
TUNNEL_ID=""
CREDENTIAL_FILE=""
HOSTNAME=""
SSO_USERNAME=""
ACCOUNT_ID=""
API_TOKEN=""
TUNNEL_TOKEN=""
INSTALL_METHOD="config" # 默认使用配置文件方式

# 检查必要工具
check_deps() {
    # 检查whiptail是否安装
    if ! command -v whiptail &> /dev/null; then
        echo -e "${YELLOW}正在安装whiptail...${NC}"
        if [ -f /etc/debian_version ]; then
            apt-get update && apt-get install -y whiptail
        elif [ -f /etc/redhat-release ]; then
            yum install -y newt
        else
            echo -e "${RED}无法安装whiptail，请手动安装后重试${NC}"
            exit 1
        fi
    fi
}

# 检查是否为root用户
check_root() {
    if [ "$EUID" -ne 0 ]; then
        echo -e "${RED}请使用root权限运行此脚本${NC}"
        exit 1
    fi
}

# 收集用户输入
collect_user_input() {
    # 显示欢迎信息
    whiptail --title "Cloudflare Tunnels SSH配置" \
             --msgbox "欢迎使用Cloudflare Tunnels SSH配置工具\n\n此脚本将帮助您设置Cloudflare Tunnels和短期证书\n以便通过Web SSH安全地访问您的服务器。" \
             15 60
    
    # 选择安装方式
    INSTALL_METHOD=$(whiptail --title "选择安装方式" \
                     --menu "请选择Cloudflare Tunnel的安装配置方式:" 15 70 2 \
                     "token" "使用Token令牌直接安装（推荐，更简单）" \
                     "config" "手动配置（高级选项，需要更多配置）" \
                     3>&1 1>&2 2>&3)
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}操作已取消${NC}"
        exit 1
    fi
    
    if [ "$INSTALL_METHOD" = "token" ]; then
        # 获取Tunnel Token
        TUNNEL_TOKEN=$(whiptail --title "Tunnel Token" \
                          --inputbox "请输入Cloudflare提供的Tunnel Token:\n(形如: eyJhIjoixxxxxxx...，在创建Tunnel后，选择要安装的连接器时提供)" \
                          10 70 \
                          3>&1 1>&2 2>&3)
        if [ $? -ne 0 ] || [ -z "$TUNNEL_TOKEN" ]; then
            echo -e "${RED}操作已取消或Token为空${NC}"
            exit 1
        fi
        
        # 获取主机名
        HOSTNAME=$(whiptail --title "主机名" \
                           --inputbox "请输入访问SSH的域名:\n(例如: terminal.mydomain.com，需要先在DNS中设置好指向Cloudflare)" \
                           10 70 \
                           3>&1 1>&2 2>&3)
        if [ $? -ne 0 ] || [ -z "$HOSTNAME" ]; then
            echo -e "${RED}操作已取消或主机名为空${NC}"
            exit 1
        fi
    else
        # 获取Tunnel ID
        TUNNEL_ID=$(whiptail --title "Tunnel ID" \
                         --inputbox "请输入您的Cloudflare Tunnel ID:\n(这是在Cloudflare Zero Trust控制台创建的Tunnel的唯一标识符)" \
                         10 60 \
                         3>&1 1>&2 2>&3)
        if [ $? -ne 0 ] || [ -z "$TUNNEL_ID" ]; then
            echo -e "${RED}操作已取消或Tunnel ID为空${NC}"
            exit 1
        fi
        
        # 获取Credential File路径
        CREDENTIAL_FILE=$(whiptail --title "凭证文件路径" \
                               --inputbox "请输入Cloudflare凭证文件的完整路径:\n\n凭证文件是创建Tunnel时由Cloudflare生成的JSON文件，\n通常命名为<tunnel-id>.json，包含连接到Cloudflare所需的密钥。\n一般保存在/root/.cloudflared/或~/.cloudflared/目录下。" \
                               15 70 "/root/.cloudflared/${TUNNEL_ID}.json" \
                               3>&1 1>&2 2>&3)
        if [ $? -ne 0 ] || [ -z "$CREDENTIAL_FILE" ]; then
            echo -e "${RED}操作已取消或凭证文件路径为空${NC}"
            exit 1
        fi
        
        # 检查凭证文件是否存在
        if [ ! -f "$CREDENTIAL_FILE" ]; then
            whiptail --title "错误" \
                     --msgbox "凭证文件 $CREDENTIAL_FILE 不存在!\n\n请确保文件路径正确。如果您尚未创建Tunnel或下载凭证文件，\n请先在Cloudflare Zero Trust控制台中完成这些步骤。" \
                     12 70
            exit 1
        fi
        
        # 获取主机名
        HOSTNAME=$(whiptail --title "主机名" \
                            --inputbox "请输入访问SSH的域名:\n(例如: terminal.mydomain.com，需要先在DNS中设置好指向Cloudflare)" \
                            10 70 \
                            3>&1 1>&2 2>&3)
        if [ $? -ne 0 ] || [ -z "$HOSTNAME" ]; then
            echo -e "${RED}操作已取消或主机名为空${NC}"
            exit 1
        fi
    fi
    
    # 获取SSO用户名
    SSO_USERNAME=$(whiptail --title "SSO用户名" \
                            --inputbox "请输入与SSO登录匹配的本地用户名:\n(重要：此用户名必须与您在SSO中使用的用户名完全一致)" \
                            10 70 \
                            3>&1 1>&2 2>&3)
    if [ $? -ne 0 ] || [ -z "$SSO_USERNAME" ]; then
        echo -e "${RED}操作已取消或用户名为空${NC}"
        exit 1
    fi
    
    # 获取Cloudflare账户ID
    ACCOUNT_ID=$(whiptail --title "Cloudflare账户ID" \
                          --inputbox "请输入您的Cloudflare账户ID:\n(可在Cloudflare仪表板的右侧栏底部找到)" \
                          10 70 \
                          3>&1 1>&2 2>&3)
    if [ $? -ne 0 ] || [ -z "$ACCOUNT_ID" ]; then
        echo -e "${RED}操作已取消或账户ID为空${NC}"
        exit 1
    fi
    
    # 获取API Token
    API_TOKEN=$(whiptail --title "Cloudflare API令牌" \
                         --inputbox "请输入您的Cloudflare API令牌:\n(需要有SSH审核编辑权限，可在Cloudflare控制台的'我的个人资料'>'API令牌'中创建)" \
                         12 70 \
                         3>&1 1>&2 2>&3)
    if [ $? -ne 0 ] || [ -z "$API_TOKEN" ]; then
        echo -e "${RED}操作已取消或API令牌为空${NC}"
        exit 1
    fi
    
    # 显示摘要并确认
    if [ "$INSTALL_METHOD" = "token" ]; then
        CONFIRM=$(whiptail --title "确认配置" --yesno "\
配置摘要:
- 安装方式: 使用Token直接安装
- Token: ${TUNNEL_TOKEN:0:10}...
- 域名: $HOSTNAME
- 用户名: $SSO_USERNAME
- 账户ID: $ACCOUNT_ID
- API令牌: ${API_TOKEN:0:5}...

是否确认继续?" 16 60 3>&1 1>&2 2>&3)
    else
        CONFIRM=$(whiptail --title "确认配置" --yesno "\
配置摘要:
- 安装方式: 手动配置
- Tunnel ID: $TUNNEL_ID
- 凭证文件: $CREDENTIAL_FILE
- 域名: $HOSTNAME
- 用户名: $SSO_USERNAME
- 账户ID: $ACCOUNT_ID
- API令牌: ${API_TOKEN:0:5}...

是否确认继续?" 18 60 3>&1 1>&2 2>&3)
    fi
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}操作已取消${NC}"
        exit 1
    fi
}

# 检查并安装cloudflared
install_cloudflared() {
    echo -e "${BLUE}步骤1: 检查并安装cloudflared...${NC}"
    
    if command -v cloudflared &> /dev/null; then
        echo -e "${GREEN}cloudflared已安装${NC}"
    else
        echo -e "${YELLOW}安装cloudflared...${NC}"
        
        # 检测操作系统并安装
        if [ -f /etc/debian_version ]; then
            # Debian/Ubuntu使用官方安装方法
            echo -e "${YELLOW}使用Cloudflare官方方法安装cloudflared...${NC}"
            
            # 添加Cloudflare GPG密钥
            sudo mkdir -p --mode=0755 /usr/share/keyrings
            curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg | sudo tee /usr/share/keyrings/cloudflare-main.gpg >/dev/null
            
            # 添加Cloudflare仓库
            echo 'deb [signed-by=/usr/share/keyrings/cloudflare-main.gpg] https://pkg.cloudflare.com/cloudflared any main' | sudo tee /etc/apt/sources.list.d/cloudflared.list
            
            # 安装cloudflared
            sudo apt-get update && sudo apt-get install -y cloudflared
        elif [ -f /etc/redhat-release ]; then
            # CentOS/RHEL
            curl -L --output cloudflared.rpm https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-x86_64.rpm
            rpm -i cloudflared.rpm
            rm cloudflared.rpm
        else
            # 其他Linux
            mkdir -p /usr/local/bin
            curl -L --output /usr/local/bin/cloudflared https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64
            chmod +x /usr/local/bin/cloudflared
        fi
        
        echo -e "${GREEN}cloudflared安装完成${NC}"
    fi
}

# 配置Cloudflare Tunnel
setup_tunnel() {
    echo -e "${BLUE}步骤2: 配置Cloudflare Tunnel...${NC}"
    
    if [ "$INSTALL_METHOD" = "token" ]; then
        # 使用Token直接安装
        echo -e "${YELLOW}使用Token安装Cloudflare Tunnel...${NC}"
        
        # 安装服务
        cloudflared service install "$TUNNEL_TOKEN"
        
        # 确保配置目录存在
        mkdir -p /etc/cloudflared
        
        # 创建配置文件以添加SSH服务
        cat > /etc/cloudflared/config.yml << EOF
# 配置由Cloudflared Token自动生成，这里仅添加SSH服务的路由
ingress:
  - hostname: $HOSTNAME
    service: ssh://localhost:22
  - service: http_status:403
EOF
        
        echo -e "${GREEN}Tunnel配置完成${NC}"
    else
        # 手动配置方式
        echo -e "${YELLOW}使用配置文件设置Cloudflare Tunnel...${NC}"
        
        # 创建或更新配置文件
        mkdir -p /etc/cloudflared
        
        # 创建配置文件
        cat > /etc/cloudflared/config.yml << EOF
tunnel: ${TUNNEL_ID}
credentials-file: ${CREDENTIAL_FILE}
ingress:
  - hostname: ${HOSTNAME}
    service: ssh://localhost:22
  - service: http_status:403
EOF
        
        echo -e "${GREEN}Tunnel配置完成${NC}"
        
        # 设置systemd服务
        if [ ! -f /etc/systemd/system/cloudflared.service ]; then
            cloudflared service install
            echo -e "${GREEN}cloudflared服务安装完成${NC}"
        fi
    fi
    
    # 重启服务
    systemctl daemon-reload
    systemctl enable cloudflared
    systemctl restart cloudflared
    
    echo -e "${GREEN}cloudflared服务已启动${NC}"
}

# 创建用户并授予权限
create_user() {
    echo -e "${BLUE}步骤3: 创建SSH用户 $SSO_USERNAME...${NC}"
    
    # 检查用户是否已存在
    if id "$SSO_USERNAME" &>/dev/null; then
        echo -e "${YELLOW}用户 $SSO_USERNAME 已存在${NC}"
    else
        # 创建用户
        useradd -m -s /bin/bash "$SSO_USERNAME"
        echo -e "${GREEN}用户 $SSO_USERNAME 创建成功${NC}"
    fi
    
    # 添加到sudo组
    if [ -f /etc/debian_version ]; then
        # Debian/Ubuntu使用sudo组
        usermod -aG sudo "$SSO_USERNAME"
        echo -e "${GREEN}用户 $SSO_USERNAME 已添加到sudo组${NC}"
    elif [ -f /etc/redhat-release ]; then
        # CentOS/RHEL使用wheel组
        usermod -aG wheel "$SSO_USERNAME"
        echo -e "${GREEN}用户 $SSO_USERNAME 已添加到wheel组${NC}"
    fi
}

# 自动配置短期证书
setup_short_lived_cert() {
    echo -e "${BLUE}步骤4: 设置Cloudflare短期证书...${NC}"
    
    # 检查是否安装了curl和jq
    if ! command -v jq &> /dev/null; then
        echo -e "${YELLOW}正在安装jq...${NC}"
        if [ -f /etc/debian_version ]; then
            apt-get update && apt-get install -y jq
        elif [ -f /etc/redhat-release ]; then
            yum install -y jq
        else
            echo -e "${RED}无法确定包管理器，请手动安装jq${NC}"
            exit 1
        fi
    fi
    
    echo -e "${YELLOW}正在从Cloudflare API获取SSH CA公钥...${NC}"
    
    # 尝试GET请求
    RESPONSE=$(curl -s "https://api.cloudflare.com/client/v4/accounts/${ACCOUNT_ID}/access/gateway_ca" \
              --header "Authorization: Bearer ${API_TOKEN}")
    
    # 检查是否成功获取
    if echo "$RESPONSE" | grep -q "access.api.error.gateway_ca_not_found"; then
        echo -e "${YELLOW}SSH CA不存在，正在创建...${NC}"
        
        # 创建新的SSH CA
        RESPONSE=$(curl -s --request POST \
                  "https://api.cloudflare.com/client/v4/accounts/${ACCOUNT_ID}/access/gateway_ca" \
                  --header "Authorization: Bearer ${API_TOKEN}")
    fi
    
    # 提取公钥
    if ! PUBLIC_KEY=$(echo "$RESPONSE" | jq -r '.result.public_key'); then
        echo -e "${RED}无法获取公钥。请检查您的ACCOUNT_ID和API_TOKEN${NC}"
        echo -e "${RED}API响应: $RESPONSE${NC}"
        exit 1
    fi
    
    if [ "$PUBLIC_KEY" = "null" ]; then
        echo -e "${RED}获取公钥失败。请检查您的ACCOUNT_ID和API_TOKEN${NC}"
        echo -e "${RED}API响应: $RESPONSE${NC}"
        exit 1
    fi
    
    echo -e "${GREEN}成功获取SSH CA公钥${NC}"
    
    # 创建ca.pub文件
    echo "$PUBLIC_KEY" > /etc/ssh/cloudflare-ca.pub
    chmod 644 /etc/ssh/cloudflare-ca.pub
    
    echo -e "${GREEN}公钥已保存至 /etc/ssh/cloudflare-ca.pub${NC}"
    
    # 更新SSH配置
    grep -q "^PubkeyAuthentication yes" /etc/ssh/sshd_config || echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config
    grep -q "^TrustedUserCAKeys /etc/ssh/cloudflare-ca.pub" /etc/ssh/sshd_config || echo "TrustedUserCAKeys /etc/ssh/cloudflare-ca.pub" >> /etc/ssh/sshd_config
    
    echo -e "${GREEN}SSH配置已更新${NC}"
    
    # 某些系统需要设置特定权限
    chmod 600 /etc/ssh/cloudflare-ca.pub
    
    # 重启SSH服务
    if systemctl is-active sshd &>/dev/null; then
        systemctl restart sshd
    elif systemctl is-active ssh &>/dev/null; then
        systemctl restart ssh
    else
        echo -e "${RED}无法确定SSH服务名，请手动重启SSH服务${NC}"
    fi
    
    echo -e "${GREEN}SSH服务已重启${NC}"
}

# 为本地用户配置SSH访问
configure_local_ssh() {
    echo -e "${BLUE}步骤5: 客户端配置指导...${NC}"
    
    echo -e "${YELLOW}在客户端使用以下命令连接到服务器:${NC}"
    echo -e "${GREEN}ssh ${SSO_USERNAME}@$(hostname -I | awk '{print $1}')${NC}"
    
    echo -e "${BLUE}用户需要通过Cloudflare WARP客户端连接到您的网络${NC}"
    echo -e "${YELLOW}确保用户在Zero Trust控制台中有适当的访问权限${NC}"
}

# 验证安装
verify_installation() {
    echo -e "${BLUE}步骤6: 验证安装...${NC}"
    
    # 检查服务状态
    if systemctl is-active cloudflared &>/dev/null; then
        echo -e "${GREEN}cloudflared服务运行正常${NC}"
    else
        echo -e "${RED}cloudflared服务未正常运行，请检查日志${NC}"
        systemctl status cloudflared
    fi
    
    # 检查SSH配置
    if grep -q "TrustedUserCAKeys /etc/ssh/cloudflare-ca.pub" /etc/ssh/sshd_config; then
        echo -e "${GREEN}SSH短期证书配置正确${NC}"
    else
        echo -e "${RED}SSH短期证书配置可能有问题${NC}"
    fi
    
    echo -e "${GREEN}设置完成!${NC}"
    echo -e "${YELLOW}请确保在Cloudflare Zero Trust控制台中配置了正确的应用和访问策略${NC}"
}

main() {
    check_root
    check_deps
    collect_user_input
    
    echo -e "${BLUE}开始配置Cloudflare Tunnels和短期证书...${NC}"
    
    install_cloudflared
    setup_tunnel
    create_user
    setup_short_lived_cert
    configure_local_ssh
    verify_installation
    
    whiptail --title "配置完成" \
             --msgbox "Cloudflare Tunnels和短期证书配置已完成!\n\n现在您可以通过 ${HOSTNAME} 使用Web SSH访问您的服务器。" \
             12 60
    
    echo -e "${GREEN}配置完成!${NC}"
    echo -e "${YELLOW}现在您可以通过 ${HOSTNAME} 使用Web SSH访问您的服务器${NC}"
}

main "$@"