#!/bin/bash

# 完全自动化SSH密钥部署脚本
# 功能：一键生成通用SSH密钥并自动部署到当前服务器
# 特性：
# 1. 完全自动化，无需任何交互
# 2. 生成通用密钥对，固定密码88888888
# 3. 自动将公钥部署到当前服务器的authorized_keys
# 4. 自动加固SSH配置(禁用密码登录)
# 5. 生成的密钥可在多台服务器上通用

# 配置参数
KEY_DIR="$HOME/.ssh"
KEY_NAME="universal_ssh_key"
KEY_PASS="88888888"
KEY_TYPE="ed25519"
SSH_PORT="22"
BACKUP_DIR="/tmp/ssh_setup_backup_$(date +%s)"

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

# 日志函数
log() {
    echo -e "$(date '+%Y-%m-%d %H:%M:%S') $1"
}

# 安装必要依赖
install_deps() {
    log "${YELLOW}检查系统依赖...${NC}"
    local pkgs=(ssh-keygen ssh sshpass expect)
    local missing=()
    
    for pkg in "${pkgs[@]}"; do
        if ! command -v "$pkg" &>/dev/null; then
            missing+=("$pkg")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        log "${YELLOW}正在安装依赖: ${missing[*]}${NC}"
        if command -v apt-get &>/dev/null; then
            apt-get update -qq && apt-get install -y -qq "${missing[@]}" >/dev/null || {
                log "${RED}依赖安装失败${NC}"
                return 1
            }
        elif command -v yum &>/dev/null; then
            yum install -y -q "${missing[@]}" >/dev/null || {
                log "${RED}依赖安装失败${NC}"
                return 1
            }
        else
            log "${RED}不支持的包管理器${NC}"
            return 1
        fi
    fi
    return 0
}

# 生成SSH密钥
generate_key() {
    log "${YELLOW}正在生成SSH密钥...${NC}"
    
    mkdir -p "$KEY_DIR" || {
        log "${RED}无法创建SSH目录${NC}"
        return 1
    }
    
    # 生成新密钥
    ssh-keygen -t "$KEY_TYPE" -C "universal_auto_key" -N "$KEY_PASS" -f "${KEY_DIR}/${KEY_NAME}" -q || {
        log "${RED}密钥生成失败${NC}"
        return 1
    }
    
    # 设置正确权限
    chmod 700 "$KEY_DIR"
    chmod 600 "${KEY_DIR}/${KEY_NAME}"
    chmod 644 "${KEY_DIR}/${KEY_NAME}.pub"
    
    log "${GREEN}SSH密钥生成成功${NC}"
    return 0
}

# 部署公钥到本机
deploy_key_locally() {
    log "${YELLOW}正在部署公钥到本机...${NC}"
    
    mkdir -p "$BACKUP_DIR"
    
    # 备份原密钥
    if [ -f "${KEY_DIR}/authorized_keys" ]; then
        cp "${KEY_DIR}/authorized_keys" "${BACKUP_DIR}/authorized_keys.bak" || {
            log "${RED}authorized_keys备份失败${NC}"
        }
    fi
    
    # 添加新公钥
    if [ -f "${KEY_DIR}/${KEY_NAME}.pub" ]; then
        mkdir -p "${KEY_DIR}"
        touch "${KEY_DIR}/authorized_keys"
        chmod 600 "${KEY_DIR}/authorized_keys"
        
        # 检查是否已存在相同公钥
        if ! grep -q "$(cat "${KEY_DIR}/${KEY_NAME}.pub")" "${KEY_DIR}/authorized_keys"; then
            cat "${KEY_DIR}/${KEY_NAME}.pub" >> "${KEY_DIR}/authorized_keys" || {
                log "${RED}公钥添加失败${NC}"
                return 1
            }
            log "${GREEN}公钥已添加到本机${NC}"
        else
            log "${YELLOW}相同公钥已存在，跳过添加${NC}"
        fi
    else
        log "${RED}公钥文件不存在${NC}"
        return 1
    fi
    
    return 0
}

# 安全加固SSH配置
secure_ssh() {
    log "${YELLOW}正在加固SSH配置...${NC}"
    
    # 备份原配置
    cp /etc/ssh/sshd_config "${BACKUP_DIR}/sshd_config.bak" || {
        log "${RED}sshd_config备份失败${NC}"
        return 1
    }
    
    # 应用安全配置
    sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
    sed -i 's/^#*PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
    sed -i 's/^#*PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
    
    # 重新加载SSH服务
    if systemctl reload sshd >/dev/null 2>&1 || service sshd reload >/dev/null 2>&1; then
        log "${GREEN}SSH配置已更新并重载${NC}"
    else
        log "${RED}SSH服务重载失败，正在回滚...${NC}"
        cp "${BACKUP_DIR}/sshd_config.bak" /etc/ssh/sshd_config
        return 1
    fi
    
    return 0
}

# 验证部署
verify_deployment() {
    log "${YELLOW}验证密钥登录...${NC}"
    
    # 测试本地密钥登录
    if ssh -i "${KEY_DIR}/${KEY_NAME}" -o PasswordAuthentication=no -o StrictHostKeyChecking=no -o ConnectTimeout=5 localhost "echo 连接测试成功" >/dev/null 2>&1; then
        log "${GREEN}密钥登录验证成功${NC}"
        return 0
    else
        log "${RED}密钥登录验证失败${NC}"
        return 1
    fi
}

# 显示成功信息
show_success() {
    echo -e "\n${GREEN}✔ SSH密钥部署成功！${NC}"
    echo -e "=============================="
    echo -e "${YELLOW}重要信息：${NC}"
    echo -e "密钥路径: ${KEY_DIR}/${KEY_NAME}"
    echo -e "密钥密码: ${KEY_PASS}"
    echo -e "\n${YELLOW}使用说明：${NC}"
    echo -e "1. 此密钥可在其他服务器上使用相同方法部署"
    echo -e "2. 连接到其他服务器时使用:"
    echo -e "   ssh -i ${KEY_DIR}/${KEY_NAME} 用户名@服务器IP"
    echo -e "3. 需要输入密钥密码时使用: ${KEY_PASS}"
    echo -e "\n${RED}⚠ 请确保在其他服务器部署前先备份现有密钥${NC}"
    echo -e "==============================\n"
}

# 主函数
main() {
    # 安装依赖
    install_deps || {
        log "${RED}依赖安装失败，退出脚本${NC}"
        exit 1
    }
    
    # 生成密钥
    generate_key || {
        log "${RED}密钥生成失败，退出脚本${NC}"
        exit 1
    }
    
    # 部署到本机
    deploy_key_locally || {
        log "${RED}本地部署失败，退出脚本${NC}"
        exit 1
    }
    
    # 安全加固
    secure_ssh || {
        log "${RED}SSH加固失败，退出脚本${NC}"
        exit 1
    }
    
    # 验证部署
    verify_deployment || {
        log "${RED}部署验证失败，退出脚本${NC}"
        exit 1
    }
    
    # 显示成功信息
    show_success
}

# 执行主函数
main