#!/bin/bash
#
# enable_root_ssh.sh — 开启 SSH 的 root 账号密码登录
#
# 特点：
#   - 幂等（重复执行不会出错，已正确配置则不做多余操作）
#   - 自动检测 sshd 配置文件位置
#   - 仅在需要修改时才修改，修改后自动重启 sshd
#   - 不碰 root 密码（密码由用户自行管理）
#
# 用法：
#   sudo ./enable_root_ssh.sh
#
# =====================================================================

set -euo pipefail

# ---- 颜色 ----
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

# ---- root 检查 ----
if [ "$(id -u)" -ne 0 ]; then
    error "请以 root 身份运行（sudo $0）"
    exit 1
fi

# ---- 检查 sshd 是否安装 ----
if ! command -v sshd &>/dev/null; then
    # 检测系统发行版，给出对应的安装命令
    error "sshd 未安装，无法继续。"
    echo
    echo "请根据你的系统类型安装 openssh-server："
    echo ""

    # 尝试读取 os-release 获取发行版信息
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID="${ID:-unknown}"
        OS_LIKE="${ID_LIKE:-}"
    else
        OS_ID="unknown"
        OS_LIKE=""
    fi

    case "$OS_ID" in
        debian|ubuntu)
            printf "  ${GREEN}Debian/Ubuntu:${NC}  apt-get update && apt-get install -y openssh-server\n"
            ;;
        rhel|centos|almalinux|rocky|fedora)
            printf "  ${GREEN}RHEL/CentOS/Alma/Rocky:${NC}  yum install -y openssh-server\n"
            printf "  ${GREEN}Fedora:${NC}               dnf install -y openssh-server\n"
            ;;
        suse|opensuse*)
            printf "  ${GREEN}SUSE/openSUSE:${NC}  zypper install -y openssh-server\n"
            ;;
        alpine)
            printf "  ${GREEN}Alpine:${NC}  apk add openssh-server\n"
            ;;
        arch|manjaro)
            printf "  ${GREEN}Arch/Manjaro:${NC}  pacman -S openssh\n"
            ;;
        *)
            # 根据 ID_LIKE 兜底判断
            case "$OS_LIKE" in
                *debian*)
                    printf "  ${GREEN}Debian/Ubuntu:${NC}  apt-get update && apt-get install -y openssh-server\n"
                    ;;
                *rhel*|*fedora*)
                    printf "  ${GREEN}RHEL/CentOS:${NC}  yum install -y openssh-server\n"
                    ;;
                *suse*)
                    printf "  ${GREEN}SUSE:${NC}  zypper install -y openssh-server\n"
                    ;;
                *)
                    printf "  ${YELLOW}无法识别系统，请手动安装 openssh-server${NC}\n"
                    printf "  常见命令：\n"
                    printf "    Debian/Ubuntu:  apt-get install -y openssh-server\n"
                    printf "    RHEL/CentOS:    yum install -y openssh-server\n"
                    printf "    Alpine:          apk add openssh-server\n"
                    printf "    Arch:            pacman -S openssh\n"
                    ;;
            esac
            ;;
    esac
    echo ""
    exit 1
fi

# ---- 定位 sshd 配置文件 ----
SSH_CONFIG=""

# 候选路径列表（覆盖常见发行版 + 编译安装路径）
CANDIDATES=(
    "/etc/ssh/sshd_config"
    "/usr/local/etc/sshd_config"
    "/etc/sshd_config"
    "/usr/etc/sshd_config"
    "/opt/etc/ssh/sshd_config"
)

for f in "${CANDIDATES[@]}"; do
    if [ -f "$f" ]; then
        SSH_CONFIG="$f"
        break
    fi
done

# 如果候选路径都没命中，试试从运行中的 sshd 进程找
if [ -z "$SSH_CONFIG" ]; then
    RUNNING_CFG=$(ps -ef | grep '[s]shd' | grep -oE '\-f\s+\S+' | awk '{print $2}' | head -1)
    if [ -n "$RUNNING_CFG" ] && [ -f "$RUNNING_CFG" ]; then
        SSH_CONFIG="$RUNNING_CFG"
    fi
fi

# 依然找不到 → 用默认路径创建
if [ -z "$SSH_CONFIG" ]; then
    SSH_CONFIG="/etc/ssh/sshd_config"
    if [ ! -f "$SSH_CONFIG" ]; then
        warn "未找到现有 sshd_config，将在 $SSH_CONFIG 新建"
        # 生成一个最小化的默认配置
        {
            echo "# $SSH_CONFIG — 由 enable_root_ssh.sh 自动生成"
            echo ""
            echo "Port 22"
            echo "AddressFamily any"
            echo "ListenAddress 0.0.0.0"
            echo "Protocol 2"
            echo "HostKey /etc/ssh/ssh_host_rsa_key"
            echo "HostKey /etc/ssh/ssh_host_ecdsa_key"
            echo "HostKey /etc/ssh/ssh_host_ed25519_key"
            echo "SyslogFacility AUTHPRIV"
            echo "LogLevel INFO"
            echo "LoginGraceTime 2m"
            echo "PermitRootLogin yes"
            echo "StrictModes yes"
            echo "MaxAuthTries 6"
            echo "MaxSessions 10"
            echo "PubkeyAuthentication yes"
            echo "PasswordAuthentication yes"
            echo "PermitEmptyPasswords no"
            echo "ChallengeResponseAuthentication no"
            echo "GSSAPIAuthentication yes"
            echo "GSSAPICleanupCredentials no"
            echo "UsePAM yes"
            echo "X11Forwarding no"
            echo "PrintMotd no"
            echo "AcceptEnv LANG LC_*"
            echo "Subsystem sftp /usr/libexec/openssh/sftp-server"
        } > "$SSH_CONFIG"
    fi
fi
if [ -z "$SSH_CONFIG" ]; then
    error "找不到 sshd 配置文件（尝试过 ${CANDIDATES[*]} 及 sshd_config.d/）"
    exit 1
fi
info "sshd 主配置文件: $SSH_CONFIG"

# 扫描所有可能包含 sshd 配置的文件（包括 sshd_config.d/）
ALL_SSH_CONFIGS=()
[ -f /etc/ssh/sshd_config ] && ALL_SSH_CONFIGS+=(/etc/ssh/sshd_config)
if [ -d /etc/ssh/sshd_config.d ]; then
    while IFS= read -r -d '' f; do
        ALL_SSH_CONFIGS+=("$f")
    done < <(find /etc/ssh/sshd_config.d -name '*.conf' -print0 2>/dev/null || true)
fi

# =====================================================================
# 函数：确保某一行配置存在于 sshd_config（幂等）
#   如果已存在（非注释行）则按需替换，否则追加到文件末尾
# =====================================================================
ensure_config() {
    local key="$1"   # e.g. PermitRootLogin
    local value="$2" # e.g. yes
    local changed=0

    # 优先处理 sshd_config 主文件
    for cfg in "${ALL_SSH_CONFIGS[@]}"; do
        # 检查该文件中是否已有此 key 的非注释行
        if grep -qE "^\s*${key}\b" "$cfg" 2>/dev/null; then
            local current
            current=$(grep -E "^\s*${key}\b" "$cfg" | head -1 | awk '{print $2}')
            if [ "$current" != "$value" ]; then
                # 替换
                sed -ri "s/^(\s*${key}\s+)\S+(.*)$/\1${value}\2/" "$cfg"
                changed=1
                info "$cfg: $key 已更新为 $value"
            fi
            # 只要某个文件里有正确的配置就行，不再处理其他文件
            return $changed
        fi
    done

    # 没有在任何文件中找到该 key 的非注释行 → 追加到主配置
    # 但先检查是否已被注释（#PermitRootLogin），有则替换注释行
    for cfg in "${ALL_SSH_CONFIGS[@]}"; do
        if grep -qE "^\s*#\s*${key}\b" "$cfg" 2>/dev/null; then
            sed -ri "s/^(\s*#\s*)(${key}\s+)\S+(.*)$/\2${value}\3/" "$cfg"
            info "$cfg: 取消注释并设置 $key $value"
            return 1
        fi
    done

    # 完全不存在 → 追加
    echo "" >> "$SSH_CONFIG"
    echo "$key $value" >> "$SSH_CONFIG"
    info "$SSH_CONFIG: 追加 $key $value"
    return 1
}

# =====================================================================
# 修改 sshd 配置
# =====================================================================
modify_sshd_config() {
    local need_restart=0

    # --- PermitRootLogin yes ---
    ensure_config "PermitRootLogin" "yes" || need_restart=1

    # --- PasswordAuthentication yes ---
    ensure_config "PasswordAuthentication" "yes" || need_restart=1

    return $need_restart
}

# =====================================================================
# 重启 sshd 服务
# =====================================================================
restart_sshd() {
    # 找到 sshd 服务名
    local svc=""
    for s in sshd ssh; do
        if systemctl list-unit-files 2>/dev/null | grep -q "^${s}\.service"; then
            svc="$s"
            break
        fi
    done

    if [ -n "$svc" ]; then
        systemctl restart "$svc"
        systemctl enable "$svc" 2>/dev/null || true
        info "已重启并启用 ${svc}.service"
    else
        # 尝试 service 命令（老系统）
        if command -v service &>/dev/null; then
            service ssh restart 2>/dev/null || service sshd restart 2>/dev/null || true
            info "已通过 service 命令重启 sshd"
        else
            # 最后手段：直接 kill -HUP
            if pgrep -x sshd &>/dev/null; then
                pkill -HUP sshd
                info "已发送 HUP 信号给 sshd 进程"
            else
                warn "找不到 sshd 服务管理器，请手动重启 sshd"
            fi
        fi
    fi
}

# =====================================================================
# main
# =====================================================================
info "开始配置 SSH root 密码登录..."

need_restart=0
modify_sshd_config || need_restart=1

if [ "$need_restart" -eq 1 ]; then
    restart_sshd
else
    info "sshd 配置已是正确状态，无需重启"
fi

# 验证
echo
echo "========== 验证配置 =========="
sshd -T 2>/dev/null | grep -E 'permitrootlogin|passwordauthentication' || \
    grep -E '^(PermitRootLogin|PasswordAuthentication)' /etc/ssh/sshd_config
echo

info "完成！root 密码登录已开启。"
info "现在可以用 'ssh root@<IP>' 登录了。"