#!/bin/bash
set -euo pipefail

# ============================================
# SSH 服务配置脚本
# 用于 安装、配置 和 启动 SSH 服务
# 用法: bash setup-ssh.sh
# 环境变量: SSH_ROOT_PASSWORD（默认 root）
# ============================================

log_info() { echo "[INFO]  $*"; }
log_warn() { echo "[WARN]  $*" >&2; }
log_error() { echo "[ERROR] $*" >&2; }

log_info "开始配置 SSH 服务..."

# -------------------------------------------
# 环境变量：SSH_ROOT_PASSWORD 可覆盖默认密码
# -------------------------------------------
ROOT_PASSWORD="${SSH_ROOT_PASSWORD:-root}"

# -------------------------------------------
# 安装 SSH 服务（仅在未安装时执行）
# -------------------------------------------
if ! command -v sshd &>/dev/null; then
	log_info "openssh-server 未安装，开始安装..."
	export DEBIAN_FRONTEND=noninteractive
	if ! apt-get update -qq 2>/dev/null; then
		log_warn "apt-get update 失败，尝试离线安装..."
	fi
	if ! apt-get install -y -qq openssh-server; then
		log_error "openssh-server 安装失败，请检查网络或镜像源"
		exit 1
	fi
	log_info "openssh-server 安装完成"
else
	log_info "openssh-server 已安装，跳过"
fi

# -------------------------------------------
# 创建 sshd 用户（如果不存在）
# -------------------------------------------
if ! id "sshd" &>/dev/null; then
	useradd -r -s /bin/false -d /var/run/sshd -c "SSH privilege separation" sshd
	log_info "已创建 sshd 用户"
fi

# -------------------------------------------
# 创建并设置运行时目录权限
# /var/run -> /run (符号链接)，只创建一次即可
# -------------------------------------------
SSHD_RUN_DIR="/run/sshd"
mkdir -p "${SSHD_RUN_DIR}"
chown root:root "${SSHD_RUN_DIR}"
chmod 755 "${SSHD_RUN_DIR}"

# -------------------------------------------
# 生成 SSH 主机密钥（如果不存在）
# 注意：在容器环境中，如需跨重启持久化，
#       请将 /etc/ssh 挂载为持久化卷
# -------------------------------------------
if [ ! -f /etc/ssh/ssh_host_rsa_key ]; then
	ssh-keygen -t rsa -b 4096 -f /etc/ssh/ssh_host_rsa_key -N ""
	chmod 600 /etc/ssh/ssh_host_rsa_key
	log_info "已生成 RSA 主机密钥"
fi
if [ ! -f /etc/ssh/ssh_host_ecdsa_key ]; then
	ssh-keygen -t ecdsa -b 521 -f /etc/ssh/ssh_host_ecdsa_key -N ""
	chmod 600 /etc/ssh/ssh_host_ecdsa_key
	log_info "已生成 ECDSA 主机密钥"
fi
if [ ! -f /etc/ssh/ssh_host_ed25519_key ]; then
	ssh-keygen -t ed25519 -f /etc/ssh/ssh_host_ed25519_key -N ""
	chmod 600 /etc/ssh/ssh_host_ed25519_key
	log_info "已生成 Ed25519 主机密钥"
fi

# -------------------------------------------
# 设置 root 用户密码
# -------------------------------------------
passwd -u root 2>/dev/null || true
echo "root:${ROOT_PASSWORD}" | chpasswd
log_info "root 密码已设置"

# -------------------------------------------
# 配置 SSH 服务——允许 root 登录 + 密码认证
# 兼顾注释和未注释两种状态
# -------------------------------------------
sed -i 's/^#*PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config
sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication yes/' /etc/ssh/sshd_config

# 检查 /etc/ssh/sshd_config.d/*.conf 中是否有覆盖配置
OVERRIDE_FILES=$(grep -rls "PermitRootLogin no\|PasswordAuthentication no" /etc/ssh/sshd_config.d/ 2>/dev/null || true)
if [ -n "${OVERRIDE_FILES}" ]; then
	log_warn "以下文件可能覆盖主配置，请手动检查:"
	for f in ${OVERRIDE_FILES}; do
		echo "       ${f}"
	done
fi

# -------------------------------------------
# 清理 ~/.ssh/config 中重复的 Host localhost 条目
# 避免端口被意外覆盖
# -------------------------------------------
SSH_CONFIG="${HOME}/.ssh/config"
if [ -f "${SSH_CONFIG}" ]; then
	HOST_COUNT=$(grep -c "^Host localhost$" "${SSH_CONFIG}" 2>/dev/null || true)
	if [ "${HOST_COUNT}" -gt 1 ]; then
		log_warn "~/.ssh/config 中发现 ${HOST_COUNT} 个重复的 Host localhost 条目，保留最后一个"
		# 用 awk 按块拆分（每块含 Host 行 + 后续缩进行），保留每个块的 Host 名
		# 只保留最后一个 Host localhost 块，其他同名块跳过
		awk '
			BEGIN { idx = 0 }
			/^Host / {
				idx++
				name[idx] = \$0
				content[idx] = ""
				next
			}
			{
				content[idx] = content[idx] \$0 "\n"
			}
			END {
				for (i = 1; i <= idx; i++) {
					if (name[i] == "Host localhost") {
						is_last = 0
						for (j = i + 1; j <= idx; j++) {
							if (name[j] == "Host localhost") is_last = 1
						}
						if (is_last) continue
					}
					print name[i]
					if (content[i] != "") printf "%s", content[i]
				}
			}
		' "${SSH_CONFIG}" >"${SSH_CONFIG}.tmp" && mv "${SSH_CONFIG}.tmp" "${SSH_CONFIG}"
		chmod 600 "${SSH_CONFIG}"
		log_info "已清理重复的 Host localhost 条目"
	fi
fi

# -------------------------------------------
# 重启 / 启动 SSH 服务
# -------------------------------------------
if service ssh status &>/dev/null; then
	service ssh restart
	log_info "SSH 服务已重启"
else
	service ssh start
	log_info "SSH 服务已启动"
fi

# -------------------------------------------
# 最终状态检查
# -------------------------------------------
sleep 1
if service ssh status >/dev/null 2>&1; then
	log_info "SSH 服务运行正常"
	echo ""
	echo "  SSH 连接信息:"
	echo "    地址: localhost"
	echo "    端口: 22"
	echo "    用户: root"
	echo "    密码: ${ROOT_PASSWORD}"
	echo ""
	echo "  提示: 可通过环境变量 SSH_ROOT_PASSWORD 自定义密码"
	echo "        例: SSH_ROOT_PASSWORD=mysecret bash setup-ssh.sh"
	echo ""
	echo "  容器外连接: ssh root@localhost -p 6***1 -A"
	echo ""
else
	echo ""
	log_error "SSH 服务启动失败"
	service ssh status 2>&1 || true
	exit 1
fi

log_info "SSH 服务配置完成"