#!/bin/bash
# ========================================================
# 🛡️  secure-ssh-setup.sh
# 自动配置 SSH 服务 | 支持密钥登录 | 防火墙放行 | 安全加固
# 作者：AI 助手 | 兼容 Ubuntu/CentOS/Debian/RHEL
# 无需 Python，一键运行
# ========================================================

set -euo pipefail  # 严格模式：出错退出、未定义变量报错、管道错误捕获

# -------------------------------
# 🔧 配置选项（可自定义）
# -------------------------------

SSH_PORT=${SSH_PORT:-22}                    # SSH 端口（默认 22）
ENABLE_PASSWORD_LOGIN=${ENABLE_PASSWORD_LOGIN:-yes}   # 是否允许密码登录
ENABLE_ROOT_LOGIN=${ENABLE_ROOT_LOGIN:-no}            # 是否允许 root 密码登录（no 更安全）
AUTO_INSTALL_SSH=${AUTO_INSTALL_SSH:-yes}             # 是否自动安装 openssh-server
USE_KEY_LOGIN=${USE_KEY_LOGIN:-yes}                   # 是否生成密钥并启用密钥登录

# -------------------------------
# 🧰 工具函数
# -------------------------------

log() {
    echo "[$(date +'%H:%M:%S')] 📝 $*"
}

success() {
    echo "[$(date +'%H:%M:%S')] ✅ $*"
}

error() {
    echo "[$(date +'%H:%M:%S')] ❌ $*" >&2
    exit 1
}

warn() {
    echo "[$(date +'%H:%M:%S')] ⚠️  $*"
}

# 检查命令是否存在
has_cmd() {
    command -v "$1" &>/dev/null
}

# 等待用户确认（可跳过）
confirm() {
    if [[ "${FORCE:-}" == "yes" ]]; then
        return 0
    fi
    read -p "$1 [y/N]: " -n 1 -r
    echo
    [[ ! $REPLY =~ ^[Yy]$ ]] && return 1
    return 0
}

# -------------------------------
# 🖥️ 检测系统信息
# -------------------------------

detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_NAME="$NAME"
        OS_ID="$ID"
        OS_VERSION="$VERSION_ID"
    elif [[ -f /etc/redhat-release ]]; then
        OS_NAME=$(cat /etc/redhat-release)
        OS_ID="rhel"
    else
        OS_NAME=$(uname -s)
        OS_ID="unknown"
    fi
    log "检测到系统: $OS_NAME"
}

# -------------------------------
# 🔧 安装 OpenSSH 服务（如未安装）
# -------------------------------

install_ssh_server() {
    if has_cmd sshd || has_cmd ssh; then
        return 0
    fi

    if [[ "$AUTO_INSTALL_SSH" != "yes" ]]; then
        error "SSH 服务未安装，且 AUTO_INSTALL_SSH=no，退出。"
    fi

    log "OpenSSH 未安装，正在安装..."

    case "$OS_ID" in
        ubuntu|debian)
            DEBIAN_FRONTEND=noninteractive apt-get update -qq
            DEBIAN_FRONTEND=noninteractive apt-get install -y openssh-server
            ;;
        centos|rhel|almalinux|rocky)
            yum install -y openssh-server || true
            # 对于较新系统使用 dnf
            if ! has_cmd sshd && has_cmd dnf; then
                dnf install -y openssh-server
            fi
            ;;
        arch)
            pacman -S openssh
            ;;
        suse|opensuse)
            zypper install -y openssh
            ;;
        *)
            error "不支持的系统类型，无法自动安装 SSH"
            ;;
    esac

    if ! has_cmd sshd && ! has_cmd ssh; then
        error "安装失败，请手动安装 openssh-server"
    fi

    success "OpenSSH 安装完成"
}

# -------------------------------
# ⚙️ 配置 SSH 服务
# -------------------------------

configure_ssh() {
    local config="/etc/ssh/sshd_config"
    local backup="${config}.backup.$(date +%s)"

    # 备份原始配置
    if [[ ! -f "$backup" && -f "$config" ]]; then
        cp "$config" "$backup"
        success "SSH 配置已备份: $backup"
    fi

    # 确保目录存在
    mkdir -p /etc/ssh

    # 写入基础配置（如果文件不存在）
    if [[ ! -f "$config" ]]; then
        log "创建新的 sshd_config 文件"
        {
            echo "Port $SSH_PORT"
            echo "PermitRootLogin $( [[ "$ENABLE_ROOT_LOGIN" == "yes" ]] && echo "yes" || echo "prohibit-password" )"
            echo "PasswordAuthentication $ENABLE_PASSWORD_LOGIN"
            echo "PubkeyAuthentication yes"
            echo "AuthorizedKeysFile .ssh/authorized_keys"
            echo "ChallengeResponseAuthentication no"
            echo "UsePAM yes"
            echo "X11Forwarding yes"
            echo "PrintMotd no"
            echo "AcceptEnv LANG LC_*"
            echo "Subsystem sftp /usr/lib/openssh/sftp-server"
        } > "$config"
    else
        # 修改现有配置
        sed -i "s/^#*Port .*/Port $SSH_PORT/" "$config"
        sed -i "s/^#*PermitRootLogin .*/PermitRootLogin $( [[ "$ENABLE_ROOT_LOGIN" == "yes" ]] && echo "yes" || echo "prohibit-password" )/" "$config"
        sed -i "s/^#*PasswordAuthentication .*/PasswordAuthentication $ENABLE_PASSWORD_LOGIN/" "$config"
        sed -i "s/^#*PubkeyAuthentication .*/PubkeyAuthentication yes/" "$config"
        sed -i "s/^#*AuthorizedKeysFile .*/AuthorizedKeysFile .ssh\/authorized_keys/" "$config"
    fi

    success "SSH 配置已更新 (端口: $SSH_PORT)"
}

# -------------------------------
# ▶️ 启动并启用 SSH 服务
# -------------------------------

start_ssh_service() {
    local service=""

    for s in sshd ssh; do
        if systemctl list-unit-files | grep -q "$s"; then
            service="$s"
            break
        fi
    done

    if [[ -z "$service" ]]; then
        error "未找到 SSH 服务，请检查是否安装正确"
    fi

    if systemctl is-active --quiet "$service"; then
        success "$service 已在运行"
    else
        log "启动 $service 服务..."
        systemctl start "$service" && systemctl enable "$service"
        success "$service 已启动并设置开机自启"
    fi

    # 检查端口是否监听
    if ! has_cmd ss && has_cmd netstat; then
        if ! netstat -tuln | grep -q ":$SSH_PORT\b"; then
            warn "端口 $SSH_PORT 未监听，尝试重启..."
            systemctl restart "$service"
            sleep 2
            if ! netstat -tuln | grep -q ":$SSH_PORT\b"; then
                error "SSH 服务启动失败，请检查日志: journalctl -u $service"
            fi
        fi
    else
        if ! ss -tuln | grep -q ":$SSH_PORT\b"; then
            warn "端口 $SSH_PORT 未监听，尝试重启..."
            systemctl restart "$service"
            sleep 2
            if ! ss -tuln | grep -q ":$SSH_PORT\b"; then
                error "SSH 服务启动失败，请检查日志: journalctl -u $service"
            fi
        fi
    fi
}

# -------------------------------
# 🔥 配置防火墙
# -------------------------------

setup_firewall() {
    if has_cmd ufw; then
        if ! ufw status | grep -q inactive; then
            ufw allow "$SSH_PORT/tcp" && success "UFW: 放行端口 $SSH_PORT"
        else
            log "UFW 处于非活动状态，启用中..."
            yes | ufw enable && ufw allow "$SSH_PORT/tcp"
            success "UFW 已启用并放行 $SSH_PORT"
        fi
    elif has_cmd firewall-cmd; then
        if firewall-cmd --state &>/dev/null; then
            firewall-cmd --permanent --add-port="$SSH_PORT"/tcp --zone=public
            firewall-cmd --reload
            success "Firewalld: 放行端口 $SSH_PORT"
        else
            warn "Firewalld 未运行"
        fi
    elif has_cmd iptables; then
        iptables -A INPUT -p tcp --dport "$SSH_PORT" -j ACCEPT
        success "Iptables: 放行端口 $SSH_PORT"
    else
        warn "未检测到 UFW、Firewalld 或 Iptables，跳过防火墙配置"
    fi
}

# -------------------------------
# 🔑 配置密钥登录
# -------------------------------

setup_key_login() {
    if [[ "$USE_KEY_LOGIN" != "yes" ]]; then
        return
    fi

    local user="${SUDO_USER:-$USER}"
    local home
    home=$(getent passwd "$user" | cut -d: -f6) || error "无法获取用户 $user 的家目录"

    if [[ -z "$home" || ! -d "$home" ]]; then
        error "用户 $user 的家目录不存在"
    fi

    local ssh_dir="$home/.ssh"
    local key_file="$ssh_dir/id_rsa"
    local pub_key="$key_file.pub"
    local auth_keys="$ssh_dir/authorized_keys"

    mkdir -p "$ssh_dir"
    chmod 700 "$ssh_dir"

    # 生成密钥（如果不存在）
    if [[ ! -f "$key_file" ]]; then
        log "为用户 $user 生成 SSH 密钥对..."
        if sudo -u "$user" ssh-keygen -t rsa -b 2048 -f "$key_file" -N "" </dev/null 2>/dev/null; then
            success "密钥已生成: $key_file"
        else
            warn "密钥生成失败，跳过"
            return
        fi
    fi

    # 添加公钥到 authorized_keys
    if [[ -f "$pub_key" ]]; then
        mkdir -p "$ssh_dir"
        cat "$pub_key" >> "$auth_keys" 2>/dev/null || true
        chmod 600 "$auth_keys"
        chown -R "$user:$user" "$ssh_dir"
        success "公钥已添加到 $auth_keys，支持密钥登录"
    fi
}

# -------------------------------
# 🧹 清理临时变量
# -------------------------------

cleanup() {
    unset SSH_PORT ENABLE_PASSWORD_LOGIN ENABLE_ROOT_LOGIN AUTO_INSTALL_SSH USE_KEY_LOGIN FORCE
}

# -------------------------------
# 🚀 主函数
# -------------------------------

main() {
    log "开始配置 SSH 服务..."

    # 提示用户
    echo
    echo "配置摘要："
    echo "  SSH 端口: $SSH_PORT"
    echo "  允许密码登录: $ENABLE_PASSWORD_LOGIN"
    echo "  允许 root 登录: $ENABLE_ROOT_LOGIN"
    echo "  自动生成密钥: $USE_KEY_LOGIN"
    echo

    if ! confirm "确认执行配置？"; then
        log "用户取消，退出。"
        exit 1
    fi

    detect_os
    install_ssh_server
    configure_ssh
    start_ssh_service
    setup_firewall
    setup_key_login

    # 最终提示
    echo
    echo "=================================================="
    echo "✅ SSH 配置完成！"
    echo "🔑 登录方式："
    echo "    密钥登录: ssh -i ~/.ssh/id_rsa $USER@localhost"
    echo "    或复制公钥到远程: ssh-copy-id user@host"
    if [[ "$SSH_PORT" != "22" ]]; then
        echo "    注意端口: -p $SSH_PORT"
    fi
    echo "💡 建议：生产环境关闭 PasswordAuthentication"
    echo "=================================================="
}

# -------------------------------
# ✅ 执行
# -------------------------------

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    if [[ $(id -u) -ne 0 ]]; then
        error "请使用 sudo 或 root 权限运行此脚本"
    fi
    main "$@"
    cleanup
fi