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

batch_mode=n

while [ "$#" -gt 0 ]; do
    case "$1" in
        --batch)
            batch_mode=y
            ;;
        -h | --help)
            printf '%s\n' \
                '用法：bash ssh [--batch]' \
                '  --batch  批量模式：添加公钥、关闭密码登录并重新加载 SSH 服务'
            exit 0
            ;;
        *)
            echo "错误：未知参数 $1" >&2
            echo "使用 --help 查看帮助" >&2
            exit 2
            ;;
    esac
    shift
done

key='ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINeLShDYnElFAKJCLH9b0uHv2rvv1QfGAZqA8RutmAw3'
expected='SHA256:Bt+kXoS7QGI014BF2ZNCu65LXx14kRT0jXWMcZem2xA'

actual="$(
    printf '%s\n' "$key" |
    ssh-keygen -lf - -E sha256 |
    awk '{print $2}'
)"

if [ "$actual" != "$expected" ]; then
    echo "错误：公钥指纹不匹配" >&2
    exit 1
fi

ssh_dir="$HOME/.ssh"
auth="$ssh_dir/authorized_keys"

install -d -m 700 "$ssh_dir"
touch "$auth"
chmod 600 "$auth"

if ! grep -qF "$key" "$auth"; then
    printf '# trah - 添加于 %s\n' "$(date '+%Y-%m-%d %H:%M:%S %z')" >> "$auth"
    printf '%s\n' "$key" >> "$auth"
fi

echo "成功：公钥已添加到 $auth"

run_as_root() {
    if [ "$(id -u)" -eq 0 ]; then
        "$@"
    else
        sudo "$@"
    fi
}

find_sshd() {
    if command -v sshd >/dev/null 2>&1; then
        command -v sshd
        return
    fi

    for path in /usr/sbin/sshd /usr/local/sbin/sshd; do
        if [ -x "$path" ]; then
            printf '%s\n' "$path"
            return
        fi
    done

    return 1
}

reload_sshd() {
    if command -v systemctl >/dev/null 2>&1; then
        if run_as_root systemctl is-active --quiet sshd 2>/dev/null; then
            run_as_root systemctl reload sshd
            return
        fi
        if run_as_root systemctl is-active --quiet ssh 2>/dev/null; then
            run_as_root systemctl reload ssh
            return
        fi
    fi

    if command -v rc-service >/dev/null 2>&1 &&
        run_as_root rc-service sshd status >/dev/null 2>&1; then
        run_as_root rc-service sshd reload
        return
    fi

    if command -v rcctl >/dev/null 2>&1 &&
        run_as_root rcctl check sshd >/dev/null 2>&1; then
        run_as_root rcctl reload sshd
        return
    fi

    if command -v service >/dev/null 2>&1; then
        if run_as_root service sshd status >/dev/null 2>&1; then
            run_as_root service sshd reload
            return
        fi
        if run_as_root service ssh status >/dev/null 2>&1; then
            run_as_root service ssh reload
            return
        fi
    fi

    return 1
}

disable_password_login() {
    local config sshd_bin candidate backup begin_marker end_marker begin_count end_count
    local effective_config pubkey_auth password_auth reload_choice

    config=/etc/ssh/sshd_config
    if [ ! -f "$config" ] && [ -f /usr/local/etc/ssh/sshd_config ]; then
        config=/usr/local/etc/ssh/sshd_config
    fi

    if [ ! -f "$config" ]; then
        echo "错误：未找到 sshd_config" >&2
        return 1
    fi

    if ! sshd_bin="$(find_sshd)"; then
        echo "错误：未找到 sshd" >&2
        return 1
    fi

    if [ "$(id -u)" -ne 0 ]; then
        if ! command -v sudo >/dev/null 2>&1; then
            echo "错误：关闭密码登录需要 root 权限或 sudo" >&2
            return 1
        fi
        run_as_root true
    fi

    begin_marker='# BEGIN trah password login policy'
    end_marker='# END trah password login policy'
    begin_count="$(run_as_root grep -cF "$begin_marker" "$config" || true)"
    end_count="$(run_as_root grep -cF "$end_marker" "$config" || true)"

    if [ "$begin_count" -gt 1 ] || [ "$end_count" -gt 1 ] ||
        [ "$begin_count" -ne "$end_count" ]; then
        echo "错误：$config 中由本脚本管理的密码登录配置块格式异常" >&2
        return 1
    fi

    candidate="$(mktemp "${TMPDIR:-/tmp}/sshd_config.XXXXXX")"
    trap "rm -f -- '$candidate'" EXIT

    awk -v begin="$begin_marker" -v end="$end_marker" '
        BEGIN {
            print begin
            print "PubkeyAuthentication yes"
            print "PasswordAuthentication no"
            print "ChallengeResponseAuthentication no"
            print end
            print ""
        }
        $0 == begin { managed = 1; next }
        $0 == end { managed = 0; next }
        !managed { print }
    ' < <(run_as_root cat "$config") > "$candidate"

    if ! run_as_root "$sshd_bin" -t -f "$candidate"; then
        echo "错误：生成的 SSH 服务配置无效，未应用任何修改" >&2
        return 1
    fi

    if ! effective_config="$(run_as_root "$sshd_bin" -T -f "$candidate")"; then
        echo "错误：无法读取生成配置的最终生效值，未应用任何修改" >&2
        return 1
    fi

    pubkey_auth="$(printf '%s\n' "$effective_config" | awk '$1 == "pubkeyauthentication" { print $2; exit }')"
    password_auth="$(printf '%s\n' "$effective_config" | awk '$1 == "passwordauthentication" { print $2; exit }')"

    if [ "$pubkey_auth" != yes ] || [ "$password_auth" != no ]; then
        echo "错误：SSH 认证配置未达到安全预期，未应用任何修改" >&2
        echo "当前检测结果：公钥认证=${pubkey_auth:-未知}，密码认证=${password_auth:-未知}" >&2
        return 1
    fi

    backup="${config}.bak.$(date '+%Y%m%d%H%M%S').$$"
    run_as_root cp -p "$config" "$backup"
    run_as_root cp "$candidate" "$config"

    if ! run_as_root "$sshd_bin" -t -f "$config"; then
        run_as_root cp -p "$backup" "$config"
        echo "错误：SSH 服务配置校验失败，已从 $backup 恢复" >&2
        return 1
    fi

    if ! effective_config="$(run_as_root "$sshd_bin" -T -f "$config")" ||
        [ "$(printf '%s\n' "$effective_config" | awk '$1 == "pubkeyauthentication" { print $2; exit }')" != yes ] ||
        [ "$(printf '%s\n' "$effective_config" | awk '$1 == "passwordauthentication" { print $2; exit }')" != no ]; then
        run_as_root cp -p "$backup" "$config"
        echo "错误：写入后的 SSH 认证配置不符合预期，已从 $backup 恢复" >&2
        return 1
    fi

    echo "成功：已开启公钥认证并关闭密码登录，原配置备份在 $backup"

    reload_choice=y
    if [ "$batch_mode" != y ] &&
        printf '是否立即重新加载 SSH 服务使配置生效？[Y/n]（默认 Y）：' 2>/dev/null > /dev/tty; then
        IFS= read -r reload_choice 2>/dev/null < /dev/tty || reload_choice=y
    fi

    if [ "$reload_choice" = n ] || [ "$reload_choice" = N ]; then
        echo "提示：尚未重新加载 SSH 服务，新配置可能暂未生效，请稍后手动重新加载"
    elif reload_sshd; then
        echo "成功：SSH 服务已重新加载，新配置已生效"
    else
        echo "警告：无法自动重新加载 SSH 服务，请手动重新加载" >&2
        echo "警告：配置文件为 $config，原配置备份在 $backup" >&2
        if [ "$batch_mode" = y ]; then
            return 1
        fi
    fi

    rm -f -- "$candidate"
    trap - EXIT
}

disable_password=n
if [ "$batch_mode" = y ]; then
    disable_password=y
    echo "提示：批量模式已启用，将关闭密码登录并重新加载 SSH 服务"
elif printf '是否关闭服务器密码登录？[y/N]（默认 N）：' 2>/dev/null > /dev/tty; then
    IFS= read -r disable_password 2>/dev/null < /dev/tty || disable_password=n
fi

case "$disable_password" in
    y | Y)
        disable_password_login
        ;;
    *)
        echo "提示：未修改服务器密码登录设置"
        ;;
esac