#!/bin/bash

# ================= 配置部分 =================
# 默认配置（可根据需要修改）
SSH_PORT=22
KEY_PATH="$HOME/.ssh/id_rsa"
REMOTE_USER="your_username"
REMOTE_HOSTS=("host1" "host2" "host3")  # 示例主机列表，可自定义

# ================= 函数定义 =================
# 彩色输出函数
color_echo() {
    local color=$1
    local message=$2
    case $color in
        red)    echo -e "\033[31m${message}\033[0m" ;;
        green)  echo -e "\033[32m${message}\033[0m" ;;
        yellow) echo -e "\033[33m${message}\033[0m" ;;
    esac
}

# 检查命令是否存在
check_command() {
    command -v "$1" >/dev/null 2>&1 || { 
        color_echo red "错误：需要安装 $1 命令"
        exit 1
    }
}

# 生成SSH密钥
generate_key() {
    if [ ! -f "${KEY_PATH}" ]; then
        color_echo yellow "正在生成新的SSH密钥..."
        ssh-keygen -t rsa -b 4096 -f "${KEY_PATH}" -N "" -C "auto-generated@$(hostname)"
        chmod 600 "${KEY_PATH}"
        chmod 700 "$HOME/.ssh"
    else
        color_echo green "使用现有SSH密钥：${KEY_PATH}"
    fi
}

# 添加公钥到远程主机
setup_ssh_auth() {
    local host=$1
    color_echo yellow "正在配置SSH认证到主机: $host"

    # 自动接受known_hosts
    ssh -o StrictHostKeyChecking=no -p "$SSH_PORT" "$REMOTE_USER@$host" "exit 2>/dev/null" || true

    # 复制公钥
    if ssh-copy-id -i "${KEY_PATH}.pub" -p "$SSH_PORT" "$REMOTE_USER@$host" >/dev/null 2>&1; then
        color_echo green "公钥已成功添加到主机: $host"
    else
        color_echo red "公钥添加失败到主机: $host"
        echo "请检查以下内容："
        echo "1. 远程服务器SSH配置是否允许密钥登录"
        echo "2. ~/.ssh 目录权限（700）"
        echo "3. ~/.ssh/authorized_keys 文件权限（600）"
        return 1
    fi
}

# 验证连接
verify_connection() {
    local host=$1
    color_echo yellow "正在验证SSH连接到主机: $host"
    if ssh -o BatchMode=yes -p "$SSH_PORT" "$REMOTE_USER@$host" "echo 'SSH验证成功'" 2>/dev/null; then
        color_echo green "SSH密钥认证配置成功！主机: $host"
        return 0
    else
        color_echo red "SSH验证失败到主机: $host"
        echo "请检查："
        echo "1. 远程服务器SSH服务是否重启"
        echo "2. /etc/ssh/sshd_config 中："
        echo "   - PubkeyAuthentication yes"
        echo "   - AuthorizedKeysFile 配置正确"
        return 1
    fi
}

# ================= 主程序 =================
color_echo green "===== SSH密钥认证批量配置脚本 ====="

# 检查必要命令
check_command ssh
check_command ssh-copy-id

# 输入参数处理
read -p "请输入远程用户名 [${REMOTE_USER}]: " input_user
REMOTE_USER=${input_user:-$REMOTE_USER}

read -p "请输入远程主机地址（空格分隔）[${REMOTE_HOSTS[*]}]: " input_hosts
REMOTE_HOSTS=(${input_hosts:-"${REMOTE_HOSTS[@]}"})

read -p "请输入SSH端口 [${SSH_PORT}]: " input_port
SSH_PORT=${input_port:-$SSH_PORT}

# 生成密钥
generate_key

# 批量处理每台主机
color_echo yellow "开始批量配置SSH密钥认证..."
success_count=0
fail_count=0
for host in "${REMOTE_HOSTS[@]}"; do
    setup_ssh_auth "$host"
    if verify_connection "$host"; then
        ((success_count++))
    else
        ((fail_count++))
    fi
done

# 汇总结果
color_echo green "===== 配置完成 ====="
echo "成功配置主机数: $success_count"
if [ "$fail_count" -gt 0 ]; then
    color_echo red "配置失败主机数: $fail_count"
fi