#!/usr/bin/env bash
set -Eeuo pipefail
IFS=$'\n\t'

PROGRAM=${0##*/}
DRY_RUN=0
ASSUME_YES=0
SKIP_KEY_INSTALL=0
BUILTIN_PUBLIC_KEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIMj0+NFHaHra1hsotLWWB1z24+fuqub2NEGilzgtElbI my-vps"
PUBLIC_KEY=""
PUBLIC_KEY_FILE=""
SSH_PORT=""
ALLOW_PORTS_RAW=""
OS_FAMILY=""
PKG_MANAGER=""
SSH_SERVICE=""
BACKUP_DIR=""
SSHD_MAIN="/etc/ssh/sshd_config"
SSHD_DROPIN_DIR="/etc/ssh/sshd_config.d"
SSHD_MANAGED_FILE="/etc/ssh/sshd_config.d/00-vps-hardening.conf"
AUTHORIZED_KEYS="/root/.ssh/authorized_keys"
MANAGED_BEGIN="# BEGIN VPS-HARDENING MANAGED BLOCK"
MANAGED_END="# END VPS-HARDENING MANAGED BLOCK"
SSH_CONFIG_CHANGED=0
SSH_CONFIG_VALIDATED=0

log() { printf '[%s] %s\n' "$1" "$2"; }
info() { log INFO "$*"; }
warn() { log WARN "$*" >&2; }
die() { log ERROR "$*" >&2; exit 1; }

on_error() {
    local line=$1 status=$2
    log ERROR "第 ${line} 行执行失败，退出码 ${status}" >&2
    if (( SSH_CONFIG_CHANGED == 1 && SSH_CONFIG_VALIDATED == 0 )); then
        warn "SSH 配置尚未通过验证；正在恢复备份。"
        restore_ssh_config || warn "自动恢复失败，请使用备份目录手动恢复：${BACKUP_DIR}"
    fi
    exit "$status"
}
trap 'on_error "$LINENO" "$?"' ERR

usage() {
    cat <<'EOF'
用法：
  vps-hardening.sh [选项]

默认使用脚本内置的 my-vps 公钥。可通过以下参数覆盖：

选项：
  --public-key-file FILE   使用指定的 OpenSSH 公钥文件
  --public-key KEY         直接提供一行 OpenSSH 公钥
  --skip-key-install       不安装内置公钥，但检查 root 已有 authorized_keys
  --ssh-port PORT          SSH TCP 端口；默认从 sshd 有效配置读取
  --allow-ports LIST       额外放行端口，逗号分隔，例如 80/tcp,443/tcp
  --dry-run                仅显示计划，不修改系统
  --yes                    跳过执行前确认
  -h, --help               显示帮助

示例：
  bash vps-hardening.sh --ssh-port 22 --allow-ports 80/tcp,443/tcp --yes
EOF
}

run() {
    if (( DRY_RUN )); then
        printf '[DRY-RUN]'
        printf ' %q' "$@"
        printf '\n'
    else
        "$@"
    fi
}

require_command() {
    command -v "$1" >/dev/null 2>&1 || die "缺少必要命令：$1"
}

parse_args() {
    while (( $# > 0 )); do
        case "$1" in
            --public-key-file)
                (( $# >= 2 )) || die "$1 缺少参数"
                PUBLIC_KEY_FILE=$2; shift 2 ;;
            --public-key)
                (( $# >= 2 )) || die "$1 缺少参数"
                PUBLIC_KEY=$2; shift 2 ;;
            --skip-key-install) SKIP_KEY_INSTALL=1; shift ;;
            --ssh-port)
                (( $# >= 2 )) || die "$1 缺少参数"
                SSH_PORT=$2; shift 2 ;;
            --allow-ports)
                (( $# >= 2 )) || die "$1 缺少参数"
                ALLOW_PORTS_RAW=$2; shift 2 ;;
            --dry-run) DRY_RUN=1; shift ;;
            --yes) ASSUME_YES=1; shift ;;
            -h|--help) usage; exit 0 ;;
            *) die "未知参数：$1" ;;
        esac
    done

    if (( SKIP_KEY_INSTALL )); then
        [[ -z "$PUBLIC_KEY" && -z "$PUBLIC_KEY_FILE" ]] || die "--skip-key-install 不能与公钥参数同时使用"
    else
        [[ -z "$PUBLIC_KEY" || -z "$PUBLIC_KEY_FILE" ]] || die "--public-key 和 --public-key-file 只能使用一个"
        if [[ -z "$PUBLIC_KEY" && -z "$PUBLIC_KEY_FILE" ]]; then
            PUBLIC_KEY=$BUILTIN_PUBLIC_KEY
        fi
    fi
}

validate_port() {
    local value=$1
    [[ "$value" =~ ^[0-9]+$ ]] || die "无效端口：$value"
    (( value >= 1 && value <= 65535 )) || die "端口超出范围：$value"
}

normalize_extra_ports() {
    EXTRA_PORTS=()
    [[ -n "$ALLOW_PORTS_RAW" ]] || return 0
    local item port protocol
    local old_ifs=$IFS
    IFS=',' read -r -a EXTRA_PORTS <<< "$ALLOW_PORTS_RAW"
    IFS=$old_ifs
    for item in "${EXTRA_PORTS[@]}"; do
        [[ "$item" =~ ^([0-9]{1,5})/(tcp|udp)$ ]] || die "无效端口格式：$item（应为 80/tcp）"
        port=${BASH_REMATCH[1]}
        protocol=${BASH_REMATCH[2]}
        validate_port "$port"
        [[ "$port/$protocol" != "$SSH_PORT/tcp" ]] || warn "$item 与 SSH 规则重复，将安全地忽略重复效果"
    done
}

detect_os() {
    [[ -r /etc/os-release ]] || die "无法读取 /etc/os-release"
    # shellcheck disable=SC1091
    source /etc/os-release
    local identity="${ID:-} ${ID_LIKE:-}"
    case "$identity" in
        *debian*|*ubuntu*)
            OS_FAMILY=debian; PKG_MANAGER=apt-get; SSH_SERVICE=ssh ;;
        *rhel*|*fedora*|*centos*|*rocky*|*almalinux*)
            OS_FAMILY=rhel
            if command -v dnf >/dev/null 2>&1; then PKG_MANAGER=dnf; else PKG_MANAGER=yum; fi
            SSH_SERVICE=sshd ;;
        *) die "不支持的发行版：${PRETTY_NAME:-${ID:-unknown}}" ;;
    esac
    info "检测到系统：${PRETTY_NAME:-$ID}（${OS_FAMILY}）"
}

find_sshd() {
    if command -v sshd >/dev/null 2>&1; then
        command -v sshd
    elif [[ -x /usr/sbin/sshd ]]; then
        printf '%s\n' /usr/sbin/sshd
    else
        die "未找到 sshd"
    fi
}

detect_ssh_port() {
    local sshd_bin=$1 detected
    if [[ -z "$SSH_PORT" ]]; then
        detected=$($sshd_bin -T 2>/dev/null | while read -r key value _; do
            if [[ "$key" == port ]]; then printf '%s\n' "$value"; break; fi
        done)
        SSH_PORT=${detected:-22}
    fi
    validate_port "$SSH_PORT"
    info "SSH 端口：${SSH_PORT}/tcp"
}

load_and_validate_public_key() {
    (( SKIP_KEY_INSTALL == 0 )) || return 0
    if [[ -n "$PUBLIC_KEY_FILE" ]]; then
        [[ -r "$PUBLIC_KEY_FILE" ]] || die "无法读取公钥文件：$PUBLIC_KEY_FILE"
        PUBLIC_KEY=$(<"$PUBLIC_KEY_FILE")
    fi
    PUBLIC_KEY=${PUBLIC_KEY//$'\r'/}
    [[ "$PUBLIC_KEY" != *$'\n'* ]] || die "公钥必须只有一行"
    [[ "$PUBLIC_KEY" != *"PRIVATE KEY"* ]] || die "检测到私钥内容；只能传入 .pub 公钥"
    [[ "$PUBLIC_KEY" =~ ^(ssh-ed25519|ssh-rsa|ecdsa-sha2-nistp(256|384|521)|sk-ssh-ed25519@openssh.com|sk-ecdsa-sha2-nistp256@openssh.com)[[:space:]]+[A-Za-z0-9+/]+={0,3}([[:space:]].*)?$ ]] || die "不是支持的 OpenSSH 单行公钥格式"

    local temp_key
    temp_key=$(mktemp)
    printf '%s\n' "$PUBLIC_KEY" > "$temp_key"
    if ! ssh-keygen -l -f "$temp_key" >/dev/null 2>&1; then
        rm -f "$temp_key"
        die "ssh-keygen 无法解析此公钥"
    fi
    rm -f "$temp_key"
}

install_root_key() {
    if (( SKIP_KEY_INSTALL )); then
        [[ -s "$AUTHORIZED_KEYS" ]] || die "$AUTHORIZED_KEYS 不存在或为空，不能安全关闭密码登录"
    else
        run install -d -m 700 -o root -g root /root/.ssh
        if (( DRY_RUN )); then
            info "将检查并幂等追加公钥到 $AUTHORIZED_KEYS"
        else
            touch "$AUTHORIZED_KEYS"
            chmod 600 "$AUTHORIZED_KEYS"
            chown root:root "$AUTHORIZED_KEYS"
            if ! grep -Fqx -- "$PUBLIC_KEY" "$AUTHORIZED_KEYS"; then
                printf '%s\n' "$PUBLIC_KEY" >> "$AUTHORIZED_KEYS"
                info "已安装 root 公钥"
            else
                info "root 公钥已存在，未重复添加"
            fi
        fi
    fi
    run chmod 700 /root/.ssh
    run chmod 600 "$AUTHORIZED_KEYS"
    run chown -R root:root /root/.ssh
    if (( DRY_RUN == 0 )); then
        [[ -s "$AUTHORIZED_KEYS" ]] || die "$AUTHORIZED_KEYS 为空"
    fi
}

create_backup() {
    BACKUP_DIR="/root/vps-hardening-backup-$(date +%Y%m%d-%H%M%S)"
    if (( DRY_RUN )); then
        info "将创建备份目录：$BACKUP_DIR"
        return 0
    fi
    install -d -m 700 "$BACKUP_DIR"
    cp -a "$SSHD_MAIN" "$BACKUP_DIR/sshd_config"
    if [[ -e "$SSHD_MANAGED_FILE" ]]; then
        cp -a "$SSHD_MANAGED_FILE" "$BACKUP_DIR/00-vps-hardening.conf"
        : > "$BACKUP_DIR/managed-file-existed"
    fi
    info "SSH 配置已备份到：$BACKUP_DIR"
}

restore_ssh_config() {
    [[ -n "$BACKUP_DIR" && -d "$BACKUP_DIR" ]] || return 1
    cp -a "$BACKUP_DIR/sshd_config" "$SSHD_MAIN"
    if [[ -f "$BACKUP_DIR/managed-file-existed" ]]; then
        cp -a "$BACKUP_DIR/00-vps-hardening.conf" "$SSHD_MANAGED_FILE"
    else
        rm -f "$SSHD_MANAGED_FILE"
    fi
    SSH_CONFIG_CHANGED=0
}

managed_ssh_content() {
    cat <<EOF
$MANAGED_BEGIN
PubkeyAuthentication yes
PasswordAuthentication no
KbdInteractiveAuthentication no
ChallengeResponseAuthentication no
PermitRootLogin prohibit-password
PermitEmptyPasswords no
MaxAuthTries 3
LoginGraceTime 30
$MANAGED_END
EOF
}

main_has_usable_dropin_include() {
    local line before_match=1 pattern
    while IFS= read -r line || [[ -n "$line" ]]; do
        [[ "$line" =~ ^[[:space:]]*# ]] && continue
        if [[ "$line" =~ ^[[:space:]]*[Mm][Aa][Tt][Cc][Hh][[:space:]] ]]; then
            before_match=0
        fi
        if (( before_match )) && [[ "$line" =~ ^[[:space:]]*[Ii][Nn][Cc][Ll][Uu][Dd][Ee][[:space:]]+(.+) ]]; then
            pattern=${BASH_REMATCH[1]}
            if [[ "$pattern" == *"/etc/ssh/sshd_config.d/"* || "$pattern" == *"sshd_config.d/"* ]]; then
                return 0
            fi
        fi
    done < "$SSHD_MAIN"
    return 1
}

reject_active_match_blocks() {
    local file line
    local files=("$SSHD_MAIN")
    if [[ -d "$SSHD_DROPIN_DIR" ]]; then
        while IFS= read -r file; do files+=("$file"); done < <(find "$SSHD_DROPIN_DIR" -maxdepth 1 -type f -name '*.conf' -print 2>/dev/null)
    fi
    for file in "${files[@]}"; do
        [[ -f "$file" ]] || continue
        while IFS= read -r line || [[ -n "$line" ]]; do
            [[ "$line" =~ ^[[:space:]]*# ]] && continue
            if [[ "$line" =~ ^[[:space:]]*[Mm][Aa][Tt][Cc][Hh][[:space:]]+ ]]; then
                die "检测到活动 SSH Match 条件块：$file。无法证明所有来源地址的 root 策略一致，拒绝自动修改"
            fi
        done < "$file"
    done
}

strip_managed_block() {
    local input=$1 output=$2
    local in_block=0 line
    : > "$output"
    while IFS= read -r line || [[ -n "$line" ]]; do
        if [[ "$line" == "$MANAGED_BEGIN" ]]; then in_block=1; continue; fi
        if [[ "$line" == "$MANAGED_END" ]]; then in_block=0; continue; fi
        (( in_block )) || printf '%s\n' "$line" >> "$output"
    done < "$input"
    (( in_block == 0 )) || die "主配置中的托管块不完整，拒绝修改"
}

write_ssh_config() {
    create_backup
    if (( DRY_RUN )); then
        if main_has_usable_dropin_include; then
            info "将写入 $SSHD_MANAGED_FILE"
        else
            info "未检测到全局 drop-in Include，将在 $SSHD_MAIN 顶部写入托管块"
        fi
        managed_ssh_content
        return 0
    fi

    local temp_main temp_conf
    temp_main=$(mktemp)
    temp_conf=$(mktemp)
    managed_ssh_content > "$temp_conf"

    if main_has_usable_dropin_include; then
        install -d -m 755 "$SSHD_DROPIN_DIR"
        install -m 644 "$temp_conf" "$SSHD_MANAGED_FILE"
        strip_managed_block "$SSHD_MAIN" "$temp_main"
        cat "$temp_main" > "$SSHD_MAIN"
    else
        strip_managed_block "$SSHD_MAIN" "$temp_main"
        { cat "$temp_conf"; printf '\n'; cat "$temp_main"; } > "$SSHD_MAIN"
        rm -f "$SSHD_MANAGED_FILE"
    fi
    rm -f "$temp_main" "$temp_conf"
    SSH_CONFIG_CHANGED=1
}

assert_effective_ssh_config() {
    local sshd_bin=$1 output key expected actual
    output=$($sshd_bin -T -C user=root,host=localhost,addr=127.0.0.1)
    while read -r key expected; do
        actual=$(while read -r found_key found_value _; do
            if [[ "$found_key" == "$key" ]]; then printf '%s\n' "$found_value"; break; fi
        done <<< "$output")
        case "$key" in
            permitrootlogin)
                if [[ "$actual" != prohibit-password && "$actual" != without-password ]]; then
                    warn "SSH 有效配置不安全：$key=${actual:-missing}"
                    return 1
                fi ;;
            *)
                if [[ "$actual" != "$expected" ]]; then
                    warn "SSH 有效配置不符合预期：$key=${actual:-missing}，期望 $expected"
                    return 1
                fi ;;
        esac
    done <<'EOF'
pubkeyauthentication yes
passwordauthentication no
kbdinteractiveauthentication no
permitrootlogin prohibit-password
permitemptypasswords no
maxauthtries 3
logingracetime 30
EOF
}

validate_and_reload_ssh() {
    local sshd_bin=$1
    if (( DRY_RUN )); then
        info "将执行：$sshd_bin -t，并检查 root 上下文的 sshd -T"
        info "验证通过后将 reload ${SSH_SERVICE}（不会 restart）"
        return 0
    fi
    if ! "$sshd_bin" -t; then
        restore_ssh_config
        die "sshd 语法检查失败，已恢复原配置且未 reload"
    fi
    if ! assert_effective_ssh_config "$sshd_bin"; then
        restore_ssh_config
        die "sshd 有效配置检查失败，已恢复原配置且未 reload"
    fi
    SSH_CONFIG_VALIDATED=1
    systemctl reload "$SSH_SERVICE"
    info "SSH 配置已验证并平滑重载"
}

install_packages() {
    if [[ "$OS_FAMILY" == debian ]]; then
        run env DEBIAN_FRONTEND=noninteractive apt-get update
        run env DEBIAN_FRONTEND=noninteractive apt-get install -y fail2ban
    else
        if ! command -v fail2ban-client >/dev/null 2>&1; then
            run "$PKG_MANAGER" install -y epel-release
        fi
        run "$PKG_MANAGER" install -y fail2ban
    fi
}

disable_firewall() {
    if [[ "$OS_FAMILY" == debian ]]; then
        if command -v ufw >/dev/null 2>&1; then
            run ufw --force disable
        else
            info "未安装 UFW，跳过"
        fi
        run systemctl disable --now ufw 2>/dev/null || true
    else
        if command -v firewall-cmd >/dev/null 2>&1; then
            run systemctl disable --now firewalld
        else
            info "未安装 firewalld，跳过"
        fi
    fi
    if (( DRY_RUN == 0 )); then
        if [[ "$OS_FAMILY" == debian ]]; then
            ufw status 2>/dev/null | grep -q '^Status: inactive$' || die "UFW 未确认已关闭"
        else
            systemctl is-active --quiet firewalld && die "firewalld 未确认已停止"
        fi
    fi
    warn "主机防火墙已关闭；公网端口将不再由 UFW/firewalld 过滤，请依赖云厂商安全组或其他网络防火墙。"
}

configure_fail2ban() {
    local jail_file=/etc/fail2ban/jail.d/sshd.local
    local content
    content=$(cat <<EOF
[sshd]
enabled = true
port = $SSH_PORT
backend = systemd
maxretry = 5
findtime = 10m
bantime = 1h
EOF
)
    if (( DRY_RUN )); then
        info "将写入 ${jail_file}："
        printf '%s\n' "$content"
        info "将验证配置并启用 Fail2ban"
        return 0
    fi
    install -d -m 755 /etc/fail2ban/jail.d
    printf '%s\n' "$content" > "$jail_file"
    chmod 644 "$jail_file"
    fail2ban-client -t
    systemctl enable fail2ban
    systemctl restart fail2ban

    local attempt
    for attempt in {1..15}; do
        if fail2ban-client ping >/dev/null 2>&1 && fail2ban-client status sshd >/dev/null 2>&1; then
            info "Fail2ban sshd jail 已启用"
            return 0
        fi
        sleep 1
    done

    warn "Fail2ban 在 15 秒内未就绪，服务状态如下："
    systemctl --no-pager --full status fail2ban >&2 || true
    warn "Fail2ban 近期日志如下："
    journalctl -u fail2ban -n 30 --no-pager >&2 || true
    die "Fail2ban 未能正常启动，请根据上方状态和日志排查"
}

confirm_plan() {
    info "即将执行：root 密钥认证、禁用 SSH 密码认证、关闭主机防火墙、启用 Fail2ban"
    info "不会配置端口放行规则；将关闭 UFW/firewalld"
    warn "请保持当前 SSH 会话在线，并确保你持有对应私钥。"
    (( ASSUME_YES || DRY_RUN )) && return 0
    [[ -t 0 ]] || die "非交互执行必须添加 --yes"
    local answer
    read -r -p "输入 yes 继续：" answer
    [[ "$answer" == yes ]] || die "用户取消"
}

print_summary() {
    cat <<EOF

完成：
  [1/4] root 公钥已检查，authorized_keys 权限已设置
  [2/4] root 密钥登录允许，密码/交互式认证已禁用
  [3/4] 主机防火墙已关闭（UFW/firewalld）
  [4/4] Fail2ban sshd jail 已启用
EOF
    [[ -n "$BACKUP_DIR" ]] && printf '  SSH 配置备份：%s\n' "$BACKUP_DIR"
    cat <<EOF

请不要立即退出当前会话。另开本地终端测试：
  ssh -p $SSH_PORT root@VPS_IP

确认密码认证失败：
  ssh -p $SSH_PORT -o PubkeyAuthentication=no -o PreferredAuthentications=password,keyboard-interactive root@VPS_IP
EOF
}

main() {
    parse_args "$@"
    (( EUID == 0 )) || die "必须以 root 运行"
    require_command ssh-keygen
    require_command systemctl
    [[ -f "$SSHD_MAIN" ]] || die "未找到 $SSHD_MAIN"

    detect_os
    local sshd_bin
    sshd_bin=$(find_sshd)
    detect_ssh_port "$sshd_bin"
    normalize_extra_ports
    load_and_validate_public_key
    reject_active_match_blocks
    confirm_plan

    install_root_key
    write_ssh_config
    validate_and_reload_ssh "$sshd_bin"
    install_packages
    disable_firewall
    configure_fail2ban
    print_summary
}

main "$@"