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

# ====== 可配置项 ======
PUB_KEY="ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQxxxxxxx your_email@example.com"
SSH_DIR="$HOME/.ssh"
AUTHORIZED_KEYS="$SSH_DIR/authorized_keys"
SSHD_CONFIG="/etc/ssh/sshd_config"

echo "[INFO] Starting server initialization..."

# ====== 1. 确保 .ssh 目录存在且权限正确 ======
echo "[INFO] Ensuring .ssh directory..."
mkdir -p "$SSH_DIR"
chmod 700 "$SSH_DIR"

# ====== 2. 添加公钥（幂等处理：避免重复）======
echo "[INFO] Adding public key..."
touch "$AUTHORIZED_KEYS"

if ! grep -qF "$PUB_KEY" "$AUTHORIZED_KEYS"; then
    echo "$PUB_KEY" >> "$AUTHORIZED_KEYS"
    echo "[INFO] Public key added."
else
    echo "[INFO] Public key already exists. Skipping."
fi

chmod 600 "$AUTHORIZED_KEYS"

# ====== 3. 修改 sshd_config ======
echo "[INFO] Configuring sshd..."

# 备份
cp "$SSHD_CONFIG" "${SSHD_CONFIG}.bak_$(date +%F_%H-%M-%S)"

# 确保 PubkeyAuthentication 启用
if grep -q "^#PubkeyAuthentication yes" "$SSHD_CONFIG"; then
    sed -i 's/^#PubkeyAuthentication yes/PubkeyAuthentication yes/' "$SSHD_CONFIG"
elif ! grep -q "^PubkeyAuthentication yes" "$SSHD_CONFIG"; then
    echo "PubkeyAuthentication yes" >> "$SSHD_CONFIG"
fi

# ====== 4. 推荐安全增强项 ======

# 禁止密码登录（互联网服务器强烈建议，内网/专网服务器可选）
# if grep -q "^#PasswordAuthentication" "$SSHD_CONFIG"; then
#     sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' "$SSHD_CONFIG"
# elif grep -q "^PasswordAuthentication" "$SSHD_CONFIG"; then
#     sed -i 's/^PasswordAuthentication.*/PasswordAuthentication no/' "$SSHD_CONFIG"
# else
#     echo "PasswordAuthentication no" >> "$SSHD_CONFIG"
# fi

# 禁止 root 直接登录（可选）
# if grep -q "^#PermitRootLogin" "$SSHD_CONFIG"; then
#     sed -i 's/^#PermitRootLogin.*/PermitRootLogin prohibit-password/' "$SSHD_CONFIG"
# elif grep -q "^PermitRootLogin" "$SSHD_CONFIG"; then
#     sed -i 's/^PermitRootLogin.*/PermitRootLogin prohibit-password/' "$SSHD_CONFIG"
# else
#     echo "PermitRootLogin prohibit-password" >> "$SSHD_CONFIG"
# fi

# ====== 5. 重启 sshd ======
echo "[INFO] Restarting sshd..."

if command -v systemctl >/dev/null 2>&1; then
    systemctl restart sshd || systemctl restart ssh
else
    service sshd restart || service ssh restart
fi

echo "[INFO] Done."