#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BACKUP_DIR="/root/security_setup_backup_$(date +%Y%m%d_%H%M%S)"
ROLLBACK_NEEDED=false
SELECTED_PUBKEYS=()

# ─────────────────────────────────────────────
# 错误回滚机制
# ─────────────────────────────────────────────
cleanup_on_error() {
    local exit_code=$?
    # EXIT 陷阱在正常退出时也会触发，仅在出错时回滚
    if [[ $exit_code -ne 0 && "$ROLLBACK_NEEDED" == "true" ]]; then
        echo ""
        echo "[ROLLBACK] 检测到错误，正在恢复备份..." >&2
        if [[ -f "$BACKUP_DIR/sshd_config" ]]; then
            cp "$BACKUP_DIR/sshd_config" /etc/ssh/sshd_config
            echo "[ROLLBACK] 已恢复 /etc/ssh/sshd_config" >&2
        fi
        if [[ -f "$BACKUP_DIR/jail.local" ]]; then
            cp "$BACKUP_DIR/jail.local" /etc/fail2ban/jail.local
            echo "[ROLLBACK] 已恢复 /etc/fail2ban/jail.local" >&2
        elif [[ -f "$BACKUP_DIR/jail.local.absent" ]]; then
            rm -f /etc/fail2ban/jail.local
            echo "[ROLLBACK] 已移除新创建的 /etc/fail2ban/jail.local" >&2
        fi
        echo "[ROLLBACK] 回滚完成。备份文件保留在: $BACKUP_DIR" >&2
    fi
}

trap cleanup_on_error ERR EXIT

# ─────────────────────────────────────────────
# 步骤 2：前置检查
# ─────────────────────────────────────────────
check_prerequisites() {
    echo "==> [1/6] 前置检查..."

    # 检查 root 权限
    if [[ $EUID -ne 0 ]]; then
        echo "[ERROR] 此脚本必须以 root 权限运行，请使用: sudo bash $0" >&2
        exit 1
    fi

    # 检查 OS
    if [[ ! -f /etc/os-release ]]; then
        echo "[ERROR] 无法检测操作系统，/etc/os-release 不存在" >&2
        exit 1
    fi
    source /etc/os-release
    if [[ "$ID" != "ubuntu" && "$ID" != "debian" && "$ID_LIKE" != *"debian"* ]]; then
        echo "[ERROR] 此脚本仅支持 Ubuntu/Debian 系统，当前系统: ${PRETTY_NAME:-未知}" >&2
        exit 1
    fi
    echo "    操作系统: ${PRETTY_NAME}"

    # 检查必要命令
    local missing=()
    for cmd in ssh-keygen systemctl apt-get; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done
    if [[ ${#missing[@]} -gt 0 ]]; then
        echo "[ERROR] 缺少必要命令: ${missing[*]}" >&2
        exit 1
    fi

    echo "    前置检查通过。"
}

# ─────────────────────────────────────────────
# 步骤 3：公钥文件发现
# ─────────────────────────────────────────────
find_pubkeys() {
    echo "==> [2/6] 扫描公钥文件..."

    local pubkey_files=()
    while IFS= read -r -d '' f; do
        pubkey_files+=("$f")
    done < <(find "$SCRIPT_DIR" -maxdepth 1 -name "*.pub" -print0 2>/dev/null | sort -z)

    if [[ ${#pubkey_files[@]} -eq 0 ]]; then
        echo "[ERROR] 未找到任何 .pub 公钥文件。" >&2
        echo "        请将公钥文件（如 mykey.pub）放置在脚本同一目录下: $SCRIPT_DIR" >&2
        exit 1
    elif [[ ${#pubkey_files[@]} -eq 1 ]]; then
        echo "    找到 1 个公钥文件: $(basename "${pubkey_files[0]}")"
        echo "    自动选择该文件。"
        SELECTED_PUBKEYS=("${pubkey_files[0]}")
    else
        echo "    找到 ${#pubkey_files[@]} 个公钥文件："
        local i=1
        for f in "${pubkey_files[@]}"; do
            printf "    [%d] %s\n" "$i" "$(basename "$f")"
            ((i++))
        done
        echo "    [A] 添加全部"
        echo ""
        while true; do
            read -rp "    请选择 [1-${#pubkey_files[@]}/A]: " choice
            choice="${choice^^}"  # 转大写
            if [[ "$choice" == "A" ]]; then
                SELECTED_PUBKEYS=("${pubkey_files[@]}")
                echo "    已选择全部 ${#pubkey_files[@]} 个公钥文件。"
                break
            elif [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= ${#pubkey_files[@]} )); then
                SELECTED_PUBKEYS=("${pubkey_files[$((choice - 1))]}")
                echo "    已选择: $(basename "${SELECTED_PUBKEYS[0]}")"
                break
            else
                echo "    无效选择，请重新输入。"
            fi
        done
    fi
}

# ─────────────────────────────────────────────
# 步骤 4：备份
# ─────────────────────────────────────────────
create_backups() {
    echo "==> [3/6] 创建备份..."
    mkdir -p "$BACKUP_DIR"

    if [[ -f /etc/ssh/sshd_config ]]; then
        cp /etc/ssh/sshd_config "$BACKUP_DIR/sshd_config"
        echo "    已备份: /etc/ssh/sshd_config"
    fi

    if [[ -f /etc/fail2ban/jail.local ]]; then
        cp /etc/fail2ban/jail.local "$BACKUP_DIR/jail.local"
        echo "    已备份: /etc/fail2ban/jail.local"
    else
        # 标记文件原本不存在，回滚时需删除
        touch "$BACKUP_DIR/jail.local.absent"
    fi

    ROLLBACK_NEEDED=true
    echo "    备份目录: $BACKUP_DIR"
}

# ─────────────────────────────────────────────
# 步骤 5：配置公钥认证
# ─────────────────────────────────────────────
configure_pubkey_auth() {
    echo "==> [4/6] 配置公钥认证..."

    local ssh_dir="/root/.ssh"
    local auth_keys="$ssh_dir/authorized_keys"

    # 确保目录和文件存在，权限正确
    mkdir -p "$ssh_dir"
    chmod 700 "$ssh_dir"
    touch "$auth_keys"
    chmod 600 "$auth_keys"

    local added=0
    local skipped=0

    for pubkey_file in "${SELECTED_PUBKEYS[@]}"; do
        echo "    处理: $(basename "$pubkey_file")"

        # 验证公钥格式
        if ! ssh-keygen -l -f "$pubkey_file" &>/dev/null; then
            echo "    [WARNING] 文件格式无效，跳过: $(basename "$pubkey_file")" >&2
            ((skipped++)) || true
            continue
        fi

        # 提取新公钥的指纹
        local new_fingerprint
        new_fingerprint="$(ssh-keygen -l -f "$pubkey_file" | awk '{print $2}')"

        # 去重检查：遍历 authorized_keys 中每行，比较指纹
        local duplicate=false
        while IFS= read -r existing_line || [[ -n "$existing_line" ]]; do
            [[ -z "$existing_line" || "$existing_line" == \#* ]] && continue
            local tmp_file
            tmp_file="$(mktemp /tmp/pubkey_check.XXXXXX)"
            echo "$existing_line" > "$tmp_file"
            local existing_fp
            existing_fp="$(ssh-keygen -l -f "$tmp_file" 2>/dev/null | awk '{print $2}')" || true
            rm -f "$tmp_file"
            if [[ "$existing_fp" == "$new_fingerprint" ]]; then
                duplicate=true
                break
            fi
        done < "$auth_keys"

        if [[ "$duplicate" == "true" ]]; then
            echo "    [SKIP] 公钥已存在（指纹: $new_fingerprint），跳过。"
            ((skipped++)) || true
        else
            cat "$pubkey_file" >> "$auth_keys"
            echo "    [OK] 已添加公钥，指纹: $new_fingerprint"
            ((added++)) || true
        fi
    done

    echo "    公钥配置完成：添加 $added 个，跳过 $skipped 个。"
}

# ─────────────────────────────────────────────
# 步骤 6：安装 fail2ban
# ─────────────────────────────────────────────
install_fail2ban() {
    echo "==> [5/6] 安装 fail2ban..."

    if command -v fail2ban-client &>/dev/null; then
        echo "    fail2ban 已安装，跳过安装步骤。"
    else
        echo "    正在安装 fail2ban..."
        apt-get update -qq
        apt-get install -y fail2ban
        echo "    fail2ban 安装完成。"
    fi

    systemctl enable fail2ban
    echo "    fail2ban 已设置为开机自启。"
}

# ─────────────────────────────────────────────
# 步骤 7：配置 fail2ban
# ─────────────────────────────────────────────
configure_fail2ban() {
    echo "==> [6/6] 配置 fail2ban..."

    cat > /etc/fail2ban/jail.local << 'EOF'
# 由 setup_server_security.sh 生成
[DEFAULT]
bantime  = -1
findtime = 10m
maxretry = 3
banaction = iptables-multiport
ignoreip = 127.0.0.1/8 ::1

[sshd]
enabled  = true
port     = ssh
filter   = sshd
logpath  = %(sshd_log)s
backend  = %(sshd_backend)s
maxretry = 3
bantime  = -1
EOF

    echo "    已写入 /etc/fail2ban/jail.local"
    echo "    参数：bantime=-1（永久封禁），maxretry=3，findtime=10m"
}

# ─────────────────────────────────────────────
# 步骤 9：服务重启与验证
# ─────────────────────────────────────────────
verify_and_report() {
    echo ""
    echo "==> 重启 fail2ban 并验证..."

    systemctl restart fail2ban

    # 验证服务状态
    local f2b_status
    f2b_status="$(systemctl is-active fail2ban 2>/dev/null || true)"
    if [[ "$f2b_status" != "active" ]]; then
        echo "[ERROR] fail2ban 服务未能正常启动（状态: $f2b_status）" >&2
        exit 1
    fi

    # 验证 jail 激活（等待服务稳定）
    sleep 2
    local jail_status
    jail_status="$(fail2ban-client status sshd 2>&1)" || true

    # 验证 authorized_keys 权限
    local ak_perm ssh_perm
    ak_perm="$(stat -c "%a" /root/.ssh/authorized_keys 2>/dev/null || echo "?")"
    ssh_perm="$(stat -c "%a" /root/.ssh 2>/dev/null || echo "?")"

    # ── 汇总报告 ──────────────────────────────
    echo ""
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║              SSH 安全配置完成 - 汇总报告                ║"
    echo "╠══════════════════════════════════════════════════════════╣"
    echo "║ [公钥认证]                                               ║"
    printf "║   %-54s ║\n" "authorized_keys 路径: /root/.ssh/authorized_keys"
    printf "║   %-54s ║\n" "目录权限: $ssh_perm  文件权限: $ak_perm"
    echo "║   已添加公钥指纹：                                       ║"
    while IFS= read -r line || [[ -n "$line" ]]; do
        [[ -z "$line" || "$line" == \#* ]] && continue
        local tmp; tmp="$(mktemp /tmp/fp_report.XXXXXX)"
        echo "$line" > "$tmp"
        local fp; fp="$(ssh-keygen -l -f "$tmp" 2>/dev/null | awk '{print $2, $4}')" || fp="(无法解析)"
        rm -f "$tmp"
        printf "║     %-52s ║\n" "$fp"
    done < /root/.ssh/authorized_keys
    echo "║                                                          ║"
    echo "║ [fail2ban]                                               ║"
    printf "║   %-54s ║\n" "服务状态: $f2b_status"
    printf "║   %-54s ║\n" "配置文件: /etc/fail2ban/jail.local"
    printf "║   %-54s ║\n" "bantime=-1（永久封禁）  maxretry=3  findtime=10m"
    echo "║                                                          ║"
    echo "║ [未修改项]                                               ║"
    printf "║   %-54s ║\n" "SSH 端口: 保持 22 不变"
    printf "║   %-54s ║\n" "PasswordAuthentication: 未禁用（保留密码登录）"
    printf "║   %-54s ║\n" "sshd_config: 未修改"
    echo "║                                                          ║"
    printf "║   %-54s ║\n" "备份位置: $BACKUP_DIR"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo ""
    echo "fail2ban-client status sshd 输出："
    echo "$jail_status"
    echo ""
    echo "手动验证命令："
    echo "  fail2ban-client get sshd bantime    # 应输出 -1"
    echo "  fail2ban-client get sshd maxretry   # 应输出 3"
    echo "  cat /root/.ssh/authorized_keys"
}

# ─────────────────────────────────────────────
# 主流程
# ─────────────────────────────────────────────
main() {
    echo ""
    echo "======================================================"
    echo "  SSH 安全配置自动化脚本"
    echo "======================================================"
    echo ""

    check_prerequisites
    find_pubkeys
    create_backups
    configure_pubkey_auth
    install_fail2ban
    configure_fail2ban
    verify_and_report

    # 正常退出，禁用回滚
    ROLLBACK_NEEDED=false
}

main "$@"