#!/bin/bash

# ========== 初始化 ========== 
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
LOG_FILE="/root/ssh_install_${TIMESTAMP}.log"
BACKUP_FILE="/root/sshd_config_${TIMESTAMP}"
PUBKEY_URL="https://api.337.plus/sh/cnfug.keys"

# ========= 日志函数 =========
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

# ========= 检测系统信息 =========
detect_os() {
    log "正在检测操作系统信息..."
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID=${ID,,}
        OS_VERSION=${VERSION_ID}
        OS_NAME=${PRETTY_NAME:-$NAME}
    elif command -v lsb_release >/dev/null 2>&1; then
        OS_NAME=$(lsb_release -d | cut -f2)
        OS_ID=$(lsb_release -i | cut -f2 | tr '[:upper:]' '[:lower:]')
        OS_VERSION=$(lsb_release -r | cut -f2)
    elif [ -f /etc/issue ]; then
        OS_NAME=$(head -n1 /etc/issue)
        OS_ID="unknown"
        OS_VERSION="unknown"
    else
        OS_NAME="unknown"
        OS_ID="unknown"
        OS_VERSION="unknown"
    fi

    log "✅ 检测到系统: $OS_NAME (ID: $OS_ID, 版本: $OS_VERSION)"
}

# ========= 安装必要软件 =========
install_packages() {
    log "正在检测 curl 和 openssh-server..."
    
    # 检查 curl 是否已安装
    if ! command -v curl >/dev/null 2>&1; then
        log "❌ 未检测到 curl，准备安装 curl..."
        case "$OS_ID" in
            ubuntu|debian|kali)
                apt-get update -y && apt-get install -y curl && log "✅ curl 安装成功"
                ;;
            centos|rocky|almalinux|fedora|rhel)
                dnf install -y curl || yum install -y curl && log "✅ curl 安装成功"
                ;;
            alpine)
                apk update && apk add curl && log "✅ curl 安装成功"
                ;;
            arch|manjaro)
                pacman -Sy --noconfirm curl && log "✅ curl 安装成功"
                ;;
            gentoo)
                emerge net-misc/curl && log "✅ curl 安装成功"
                ;;
            opensuse*|sles)
                zypper install -y curl && log "✅ curl 安装成功"
                ;;
            openbsd)
                log "✅ OpenBSD 通常已内置 curl，跳过安装。"
                ;;
            *)
                log "❌ 无法自动识别的系统，请手动安装 curl"
                ;;
        esac
    else
        log "✅ curl 已安装，跳过安装。"
    fi

    # 检查是否已安装 openssh-server
    if [ ! -f /etc/ssh/sshd_config ]; then
        log "❌ 未检测到 /etc/ssh/sshd_config，准备安装 openssh-server..."

        case "$OS_ID" in
            ubuntu|debian|kali)
                apt-get update -y && apt-get install -y openssh-server && log "✅ openssh-server 安装成功"
                ;;
            centos|rocky|almalinux|fedora|rhel)
                dnf install -y openssh-server || yum install -y openssh-server && log "✅ openssh-server 安装成功"
                ;;
            alpine)
                apk update && apk add openssh && log "✅ openssh 安装成功"
                ;;
            arch|manjaro)
                pacman -Sy --noconfirm openssh && log "✅ openssh 安装成功"
                ;;
            gentoo)
                emerge net-misc/openssh && log "✅ openssh 安装成功"
                ;;
            opensuse*|sles)
                zypper install -y openssh && log "✅ openssh 安装成功"
                ;;
            openbsd)
                log "✅ OpenBSD 通常已内置 OpenSSH，跳过安装。"
                ;;
            *)
                log "❌ 无法自动识别的系统，请手动安装 openssh-server"
                ;;
        esac
    else
        log "✅ 已检测到 openssh-server，跳过安装。"
    fi
}
# ========= 配置 SSH安全设置 =========
configure_sshd() {
    SSHD_CONFIG="/etc/ssh/sshd_config"

    if [ ! -f "$SSHD_CONFIG" ]; then
        log "❌ $SSHD_CONFIG 不存在，无法配置。"
        exit 1
    fi

    cp "$SSHD_CONFIG" "$BACKUP_FILE"
    log "✅ 已备份 SSH 配置文件至: $BACKUP_FILE"

    log "✅ 开始配置 SSH 安全设置..."

    modify_or_append() {
        local key="$1"
        local value="$2"
        local config="$3"
        if grep -Ei "^[#\s]*$key\s+" "$config" >/dev/null; then
            local original=$(grep -Ei "^[#\s]*$key\s+.*" "$config")
            sed -i -E "s|^[#\s]*$key\s+.*|$key $value|i" "$config" && log "✅ 已修改 $key：从 -> '$original' 修改为 -> '$key $value'"
        else
            echo "$key $value" >> "$config" && log "✅ 未找到 $key，已追加为 -> '$key $value'"
        fi
    }

    modify_or_append "PermitRootLogin" "prohibit-password" "$SSHD_CONFIG"
    modify_or_append "PasswordAuthentication" "no" "$SSHD_CONFIG"
    modify_or_append "PubkeyAuthentication" "yes" "$SSHD_CONFIG"
    modify_or_append "ChallengeResponseAuthentication" "no" "$SSHD_CONFIG"
    modify_or_append "PermitEmptyPasswords" "no" "$SSHD_CONFIG"
    modify_or_append "X11Forwarding" "no" "$SSHD_CONFIG"
    modify_or_append "MaxAuthTries" "3" "$SSHD_CONFIG"
    modify_or_append "LogLevel" "VERBOSE" "$SSHD_CONFIG"

    if grep -q -r '^[#\s]*PasswordAuthentication\s\+yes' /etc/ssh/sshd_config.d/ 2>/dev/null; then
        log "✅ 检测到 /etc/ssh/sshd_config.d/ 中启用了密码登录，正在强制禁用..."
        find /etc/ssh/sshd_config.d/ -type f -exec sed -i 's/^[#\s]*PasswordAuthentication\s\+yes/PasswordAuthentication no/' {} \;
        log "✅ 已禁用 sshd_config.d 中的密码登录设置。"
    else
        log "✅ sshd_config.d 中未检测到开启密码登录的设置。"
    fi
}

# ========= 公钥设置 =========
setup_pubkey() {
    log "✅ 正在设置公钥认证..."
    mkdir -p /root/.ssh
    curl -fsSL "$PUBKEY_URL" -o /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys
    chmod 700 /root/.ssh && log "✅ 已设置 /root/.ssh/authorized_keys"
}

#restart_ssh_service() {
#    log "✅ 尝试重启 SSH 服务..."
#    
#    # 判断是否有 systemctl
#    if command -v systemctl >/dev/null 2>&1; then
#        # 优先 systemd
#        if systemctl list-units | grep -q 'sshd.service'; then
#            systemctl restart sshd && log "✅ 已通过 systemctl 重启 sshd 服务。"
#        elif systemctl list-units | grep -q 'ssh.service'; then
#            systemctl restart ssh && log "✅ 已通过 systemctl 重启 ssh 服务。"
#        else
#            log "❌ 未找到有效的 ssh 服务。"
#        fi
#    elif command -v service >/dev/null 2>&1; then
#        service sshd restart && log "✅ 已通过 service 重启 sshd 服务。"
#    elif command -v rc-service >/dev/null 2>&1; then
#        rc-service sshd restart && log "✅ 已通过 rc-service 重启 sshd 服务。"
#    elif command -v sv >/dev/null 2>&1; then
#        sv restart sshd && log "✅ 已通过 runit sv 重启 sshd。"
#    else
#        log "❌ 无法识别的 SSH 服务管理方式，请手动重启 SSH。"
#    fi
#}

restart_ssh_service() {
    log "✅ 尝试重启 SSH 服务..."

    # 优先检查 systemctl
    if command -v systemctl >/dev/null 2>&1; then
        if systemctl is-active --quiet sshd || systemctl is-active --quiet ssh; then
            systemctl restart sshd && log "✅ 已通过 systemctl 重启 sshd 服务。"
        else
            log "❌ 未找到有效的 SSH 服务。"
        fi
    # 如果没有 systemctl，检查 service
    elif command -v service >/dev/null 2>&1; then
        service sshd restart && log "✅ 已通过 service 重启 sshd 服务。"
    # 如果没有 service，检查 rc-service（Alpine Linux）
    elif command -v rc-service >/dev/null 2>&1; then
        rc-service sshd restart && log "✅ 已通过 rc-service 重启 sshd 服务。"
    # 如果没有 rc-service，检查 runit 的 sv
    elif command -v sv >/dev/null 2>&1; then
        sv restart sshd && log "✅ 已通过 runit sv 重启 sshd。"
    else
        log "❌ 无法识别的 SSH 服务管理方式，请手动重启 SSH。"
    fi
}


# ========= 主执行流程 =========
main() {
    detect_os
    install_packages
    configure_sshd
    setup_pubkey
    restart_ssh_service
    log "✅ SSH 配置与服务重启完成，日志文件: $LOG_FILE"
}

main