#!/bin/bash
# 多节点互相同步SSH密钥脚本（支持从文件读取IP、指定用户和密码）

# ==================== 颜色与样式定义 ====================
GREEN="\033[32m"    # 成功/完成
YELLOW="\033[33m"   # 进行中/提示
RED="\033[31m"      # 错误/失败
BLUE="\033[34m"     # 标题/分隔线
CYAN="\033[36m"     # 信息/说明
NC="\033[0m"        # 重置颜色

# 打印分隔线
print_sep() {
    echo -e "${BLUE}------------------------------------------------${NC}"
}

# ==================== 参数解析与校验 ====================
# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        --file=*)
            IP_FILE="${1#*=}"
            shift
            ;;
        --user=*)
            SSH_USER="${1#*=}"
            shift
            ;;
        --passwd=*)
            SSH_PASSWD="${1#*=}"
            shift
            ;;
        *)
            echo -e "${RED}未知参数: $1${NC}"
            exit 1
            ;;
    esac
done

# 检查必要参数
if [ -z "$IP_FILE" ] || [ -z "$SSH_USER" ] || [ -z "$SSH_PASSWD" ]; then
    echo -e "\n${YELLOW}用法:${NC} --file=IP文件路径 --user=用户名 --passwd=密码"
    exit 1
fi

# 检查IP文件是否存在
if [ ! -f "$IP_FILE" ]; then
    echo -e "${RED}错误：IP文件 $IP_FILE 不存在${NC}"
    exit 1
fi

# 从文件读取IP列表（去重、过滤空行）
NODES=()
while IFS= read -r ip; do
    if [ -n "$ip" ] && ! [[ " ${NODES[@]} " =~ " $ip " ]]; then
        NODES+=("$ip")
    fi
done < "$IP_FILE"

if [ ${#NODES[@]} -eq 0 ]; then
    echo -e "${RED}错误：IP文件 $IP_FILE 中未找到有效IP${NC}"
    exit 1
fi

# ==================== 核心功能函数 ====================
# 安装sshpass（处理自动输入密码）
install_sshpass() {
    echo -e "\n${YELLOW}【步骤1/3】检查并安装sshpass工具...${NC}"
    if ! command -v sshpass &> /dev/null; then
        echo -e "  ${YELLOW}→ sshpass未安装，正在安装...${NC}"
        if command -v apt &> /dev/null; then
            apt update -y >/dev/null && apt install -y sshpass >/dev/null
        elif command -v yum &> /dev/null; then
            yum install -y sshpass >/dev/null
        else
            echo -e "  ${RED}✗ 未找到apt或yum，无法安装sshpass${NC}"
            exit 1
        fi
        # 验证安装结果
        if command -v sshpass &> /dev/null; then
            echo -e "  ${GREEN}✓ sshpass安装成功${NC}"
        else
            echo -e "  ${RED}✗ sshpass安装失败，请手动安装后重试${NC}"
            exit 1
        fi
    else
        echo -e "  ${GREEN}✓ sshpass已安装${NC}"
    fi
}

# 在目标节点生成SSH密钥（若不存在）
generate_key() {
    local node=$1
    echo -e "  ${YELLOW}→ 处理节点 $node...${NC}"
    # 远程生成密钥（静默模式）
    sshpass -p "$SSH_PASSWD" ssh -o StrictHostKeyChecking=no ${SSH_USER}@$node << EOF >/dev/null 2>&1
        if [ ! -f ~/.ssh/id_rsa ]; then
            ssh-keygen -t rsa -N "" -f ~/.ssh/id_rsa
        fi
EOF
    # 检查执行结果
    if [ $? -eq 0 ]; then
        echo -e "  ${GREEN}✓ 密钥生成完成（已存在则跳过）${NC}"
    else
        echo -e "  ${RED}✗ 生成失败（可能密码错误或节点不可达）${NC}"
        exit 1
    fi
}

# 获取目标节点的公钥
get_public_key() {
    local node=$1
    # 静默获取公钥，忽略错误输出
    sshpass -p "$SSH_PASSWD" ssh -o StrictHostKeyChecking=no ${SSH_USER}@$node "cat ~/.ssh/id_rsa.pub" 2>/dev/null
}

# 将公钥追加到目标节点的authorized_keys
append_public_key() {
    local pub_key="$1"
    local target_node=$2
    # 远程执行追加公钥操作
    sshpass -p "$SSH_PASSWD" ssh -o StrictHostKeyChecking=no ${SSH_USER}@$target_node << EOF >/dev/null 2>&1
        mkdir -p ~/.ssh
        chmod 700 ~/.ssh
        if ! grep -q "$pub_key" ~/.ssh/authorized_keys 2>/dev/null; then
            echo "$pub_key" >> ~/.ssh/authorized_keys
            chmod 600 ~/.ssh/authorized_keys
        fi
EOF
    # 反馈结果
    if [ $? -eq 0 ]; then
        echo -e "    ${GREEN}→ 同步成功${NC}"
    else
        echo -e "    ${RED}→ 同步失败${NC}"
        exit 1
    fi
}

# ==================== 主流程 ====================
main() {
    # 欢迎信息
    clear
    print_sep
    echo -e "${BLUE}          多节点SSH密钥互相同步工具          ${NC}"
    print_sep
    echo -e "  ${CYAN}节点列表:${NC} ${NODES[*]}"
    echo -e "  ${CYAN}操作用户:${NC} $SSH_USER"
    echo -e "  ${CYAN}开始时间:${NC} $(date +'%Y-%m-%d %H:%M:%S')"
    print_sep

    # 1. 安装依赖工具
    install_sshpass
    print_sep

    # 2. 为所有节点生成密钥
    node_count=${#NODES[@]}
    echo -e "\n${YELLOW}【步骤2/3】为所有节点生成SSH密钥（共 $node_count 个节点）...${NC}"
    for i in "${!NODES[@]}"; do
        echo -e "\n  节点 ${i+1}/$node_count: ${NODES[$i]}"
        generate_key ${NODES[$i]}
    done
    print_sep

    # 3. 收集公钥并互相同步
    echo -e "\n${YELLOW}【步骤3/3】公钥互相同步（共 $node_count 个节点）...${NC}"
    for i in "${!NODES[@]}"; do
        current_node=${NODES[$i]}
        echo -e "\n  处理第 ${i+1}/$node_count 个节点: $current_node"
        
        # 获取当前节点公钥
        echo -e "  ${YELLOW}→ 获取公钥...${NC}"
        pub_key=$(get_public_key $current_node)
        if [ -z "$pub_key" ]; then
            echo -e "  ${RED}✗ 无法获取 $current_node 的公钥${NC}"
            exit 1
        fi
        
        # 同步到所有节点
        echo -e "  ${YELLOW}→ 同步到所有节点...${NC}"
        for target in "${NODES[@]}"; do
            echo -e "    目标节点: $target"
            append_public_key "$pub_key" $target
        done
    done

    # 完成提示
    print_sep
    echo -e "\n${GREEN}=== 所有操作完成！===${NC}"
    echo -e "  ${CYAN}验证方法:${NC} 在任意节点执行 'ssh ${NODES[1]}' 测试免密登录"
    echo -e "  ${CYAN}结束时间:${NC} $(date +'%Y-%m-%d %H:%M:%S')"
    print_sep
}

# 执行主流程
main