#!/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=""
CA_PUB_KEY=""
TUNNEL_TOKEN=""
INSTALL_METHOD="config" # 默认使用配置文件方式

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

# 打印帮助信息
print_usage() {
    echo -e "${BLUE}Cloudflare Tunnels SSH配置工具${NC}"
    echo -e "此脚本将帮助您设置Cloudflare Tunnels和短期证书以便通过Web SSH安全地访问您的服务器。"
    echo
    echo -e "使用方法:"
    echo -e "  ${YELLOW}./setup-cloudflare-ssh-access.sh${NC}"
    echo
    echo -e "您将需要提供以下信息:"
    echo -e "  1. 安装方式: 'token'(推荐) 或 'config'"
    echo -e "  2. 对于token方式: Cloudflare Tunnel Token"
    echo -e "  3. 对于config方式: Tunnel ID 和凭证文件路径"
    echo -e "  4. 访问SSH的域名 (例如: terminal.mydomain.com)"
    echo -e "  5. 与SSO登录匹配的本地用户名"
    echo -e "  6. Cloudflare SSH证书公钥(从Zero Trust控制台获取)"
    echo
}

# 读取用户输入
read_input() {
    echo -e "${GREEN}=====================================================${NC}"
    echo -e "${BLUE}欢迎使用Cloudflare Tunnels SSH配置工具${NC}"
    echo -e "${GREEN}=====================================================${NC}"
    echo
    
    # 选择安装方式
    echo -e "${YELLOW}请选择Cloudflare Tunnel的安装配置方式:${NC}"
    echo -e "1. 使用Token令牌直接安装（推荐，更简单）"
    echo -e "2. 手动配置（高级选项，需要更多配置）"
    read -p "请输入选择 [1/2]: " install_choice
    
    if [ "$install_choice" = "1" ]; then
        INSTALL_METHOD="token"
        
        # 获取Tunnel Token
        echo -e "${YELLOW}请输入Cloudflare提供的Tunnel Token:${NC}"
        echo -e "(形如: eyJhIjoixxxxxxx...，在创建Tunnel后，选择要安装的连接器时提供)"
        read -p "> " TUNNEL_TOKEN
        if [ -z "$TUNNEL_TOKEN" ]; then
            echo -e "${RED}错误: Token不能为空${NC}"
            exit 1
        fi
        
        # 获取主机名
        echo -e "${YELLOW}请输入访问SSH的域名:${NC}"
        echo -e "(例如: terminal.mydomain.com，需要先在DNS中设置好指向Cloudflare)"
        read -p "> " HOSTNAME
        if [ -z "$HOSTNAME" ]; then
            echo -e "${RED}错误: 域名不能为空${NC}"
            exit 1
        fi
    else
        INSTALL_METHOD="config"
        
        # 获取Tunnel ID
        echo -e "${YELLOW}请输入您的Cloudflare Tunnel ID:${NC}"
        echo -e "(这是在Cloudflare Zero Trust控制台创建的Tunnel的唯一标识符)"
        read -p "> " TUNNEL_ID
        if [ -z "$TUNNEL_ID" ]; then
            echo -e "${RED}错误: Tunnel ID不能为空${NC}"
            exit 1
        fi
        
        # 获取Credential File路径
        echo -e "${YELLOW}请输入Cloudflare凭证文件的完整路径:${NC}"
        echo -e "凭证文件是创建Tunnel时由Cloudflare生成的JSON文件，"
        echo -e "通常命名为<tunnel-id>.json，包含连接到Cloudflare所需的密钥。"
        echo -e "一般保存在/root/.cloudflared/或~/.cloudflared/目录下。"
        read -p "> " CREDENTIAL_FILE
        if [ -z "$CREDENTIAL_FILE" ]; then
            echo -e "${RED}错误: 凭证文件路径不能为空${NC}"
            exit 1
        fi
        
        # 检查凭证文件是否存在
        if [ ! -f "$CREDENTIAL_FILE" ]; then
            echo -e "${RED}错误: 凭证文件 $CREDENTIAL_FILE 不存在!${NC}"
            echo -e "${RED}请确保文件路径正确。如果您尚未创建Tunnel或下载凭证文件，${NC}"
            echo -e "${RED}请先在Cloudflare Zero Trust控制台中完成这些步骤。${NC}"
            exit 1
        fi
        
        # 获取主机名
        echo -e "${YELLOW}请输入访问SSH的域名:${NC}"
        echo -e "(例如: terminal.mydomain.com，需要先在DNS中设置好指向Cloudflare)"
        read -p "> " HOSTNAME
        if [ -z "$HOSTNAME" ]; then
            echo -e "${RED}错误: 域名不能为空${NC}"
            exit 1
        fi
    fi
    
    # 获取SSO用户名
    echo -e "${YELLOW}请输入与SSO登录匹配的本地用户名:${NC}"
    echo -e "(重要：此用户名必须与您在SSO中使用的用户名完全一致)"
    read -p "> " SSO_USERNAME
    if [ -z "$SSO_USERNAME" ]; then
        echo -e "${RED}错误: 用户名不能为空${NC}"
        exit 1
    fi
    
    # 获取SSH CA公钥
    echo -e "${YELLOW}请输入Cloudflare SSH证书公钥:${NC}"
    echo -e "在Zero Trust控制台 > Access > Service Auth > SSH 中找到并复制"
    echo -e "格式如: ecdsa-sha2-nistp256 AAAA... open-ssh-ca@cloudflareaccess.org"
    read -p "> " CA_PUB_KEY
    if [ -z "$CA_PUB_KEY" ]; then
        echo -e "${RED}错误: 证书公钥不能为空${NC}"
        exit 1
    fi
    
    # 验证公钥格式
    if [[ ! "$CA_PUB_KEY" =~ ^ecdsa-sha2-nistp256[[:space:]] ]]; then
        echo -e "${YELLOW}警告: 公钥格式似乎不正确，应以'ecdsa-sha2-nistp256'开头${NC}"
        echo -e "${YELLOW}确定要继续吗? [y/N]${NC}"
        read -p "> " confirm_key
        if [[ ! "$confirm_key" =~ ^[Yy]$ ]]; then
            echo -e "${RED}操作已取消${NC}"
            exit 1
        fi
    fi
    
    # 显示摘要并确认
    echo -e "${GREEN}=====================================================${NC}"
    echo -e "${BLUE}配置摘要:${NC}"
    
    if [ "$INSTALL_METHOD" = "token" ]; then
        echo -e "安装方式: 使用Token直接安装"
        echo -e "Token: ${TUNNEL_TOKEN:0:10}..."
    else
        echo -e "安装方式: 手动配置"
        echo -e "Tunnel ID: $TUNNEL_ID"
        echo -e "凭证文件: $CREDENTIAL_FILE"
    fi
    
    echo -e "域名: $HOSTNAME"
    echo -e "用户名: $SSO_USERNAME"
    echo -e "SSH证书公钥: ${CA_PUB_KEY:0:30}..."
    echo -e "${GREEN}=====================================================${NC}"
    
    echo -e "${YELLOW}是否确认继续? [y/N]${NC}"
    read -p "> " confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; 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}"
    
    echo -e "${YELLOW}正在配置SSH使用Cloudflare短期证书...${NC}"
    
    # 将公钥保存到SSH配置目录
    echo "$CA_PUB_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
    print_usage
    read_input
    
    echo -e "${BLUE}开始配置Cloudflare Tunnels和短期证书...${NC}"
    
    install_cloudflared
    setup_tunnel
    create_user
    setup_short_lived_cert
    configure_local_ssh
    verify_installation
    
    echo -e "${GREEN}=====================================================${NC}"
    echo -e "${GREEN}配置完成!${NC}"
    echo -e "${GREEN}现在您可以通过 ${HOSTNAME} 使用Web SSH访问您的服务器${NC}"
    echo -e "${GREEN}=====================================================${NC}"
}

main "$@"