#!/bin/bash
#
# ubuntu-bootstrap.sh — Ubuntu 新装系统一键初始化脚本
#
# 功能：
#   1. 更新系统软件包
#   2. 安装常用基础工具
#   3. 安装并启用 OpenSSH Server，允许 SSH 远程登录
#   4. 配置 SSH 公钥免密登录 & 防火墙 & fail2ban
#   5. 设置时区为中国时区
#   6. 安装 Tailscale 客户端并注册到自托管 Headscale 控制平面
#
# 用法：
#   curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo bash
#
#   或下载后执行：
#   curl -fsSL -o /tmp/bootstrap.sh https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh
#   sudo bash /tmp/bootstrap.sh
#
#   指定 SSH 公钥（免密登录）：
#   curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo PUBKEY="ssh-ed25519 AAAA..." bash
#
#   Tailscale 已内置默认 preauth key（reusable, 10年有效），无需手动传参即可自动注册：
#   curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo bash
#
#   如需覆盖默认 key（例如使用一次性 key）：
#   curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo TS_AUTHKEY="新的key" bash
#
#   指定时区：
#   curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo TZ="Asia/Shanghai" bash
#

set -euo pipefail

# ─── 颜色输出 ───────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

info()    { echo -e "${GREEN}[INFO]${NC}  $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error()   { echo -e "${RED}[ERROR]${NC} $*"; }
fatal()   { echo -e "${RED}[FATAL]${NC} $*"; exit 1; }
step()    { echo -e "\n${BLUE}═══════════════════════════════════════════════════════════${NC}"; echo -e "${BLUE}  $*${NC}"; echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}"; }

# ─── 前置检查 ───────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    fatal "此脚本需要 root 权限，请使用 sudo 运行: sudo bash $0"
fi

if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    if [[ "$ID" != "ubuntu" ]] && [[ "$ID_LIKE" != *"ubuntu"* ]] && [[ "$ID" != "debian" ]]; then
        warn "检测到系统为 ${PRETTY_NAME:-unknown}，本脚本针对 Ubuntu/Debian 设计，可能不完全兼容。"
        read -rp "是否继续？(y/N) " yn
        [[ "$yn" =~ ^[Yy]$ ]] || exit 0
    fi
else
    fatal "无法检测操作系统（缺少 /etc/os-release），退出。"
fi

# ─── 默认参数 ───────────────────────────────────────────
PUBKEY="${PUBKEY:-}"
TZ_SETTING="${TZ:-Asia/Shanghai}"
SSH_PORT="${SSH_PORT:-22}"
NEW_USER="${NEW_USER:-}"
SUDO_NOPASSWD="${SUDO_NOPASSWD:-true}"

# Headscale 控制平面（自托管，namespace: wendao）
HEADSCALE_URL="${HEADSCALE_URL:-https://headscale.fanmikeji.cn}"
TS_AUTHKEY="${TS_AUTHKEY:-72d4076697966cc9b538d911d1c3cbe709bb072ee5c843c8}"
TS_EXIT_NODE="${TS_EXIT_NODE:-false}"
TS_HOSTNAME="${TS_HOSTNAME:-}"

echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║         Ubuntu 新装系统一键初始化脚本                        ║${NC}"
echo -e "${GREEN}║         bootstrap.tun.fanmikeji.cn                           ║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo -e "  系统:     ${PRETTY_NAME:-unknown}"
echo -e "  时区:     ${TZ_SETTING}"
echo -e "  SSH端口:  ${SSH_PORT}"
echo -e "  Headscale: ${HEADSCALE_URL}"
[[ -n "$PUBKEY" ]]       && echo -e "  公钥:     已通过环境变量指定" || echo -e "  公钥:     未指定（跳过免密配置）"
[[ -n "$NEW_USER" ]]     && echo -e "  新用户:   ${NEW_USER}"
[[ "$SUDO_NOPASSWD" == "true" ]] && echo -e "  免密sudo:  已开启"
[[ -n "$TS_AUTHKEY" ]]   && echo -e "  TS Key:   已提供，将自动注册" || echo -e "  TS Key:   未提供（安装后手动注册）"
[[ "$TS_EXIT_NODE" == "true" ]] && echo -e "  Exit Node: 是"
[[ -n "$TS_HOSTNAME" ]]  && echo -e "  主机名:   ${TS_HOSTNAME}"
echo ""

# ─── Step 1: 更新系统 ──────────────────────────────────
step "Step 1/6: 更新系统软件包"
export DEBIAN_FRONTEND=noninteractive
apt-get update -y
apt-get upgrade -y
apt-get dist-upgrade -y
info "系统更新完成"

# ─── Step 2: 安装基础工具 ─────────────────────────────
step "Step 2/6: 安装基础软件工具"

BASIC_PACKAGES=(
    curl wget git vim nano
    htop tree unzip zip gzip bzip2 xz-utils
    net-tools iputils-ping dnsutils traceroute
    telnet ncdu jq
    software-properties-common apt-transport-https
    ca-certificates gnupg lsb-release
    build-essential make
    rsync tmux screen
    lsof strace
    python3 python3-pip
    ufw fail2ban
)

info "安装 ${#BASIC_PACKAGES[@]} 个基础软件包..."
apt-get install -y "${BASIC_PACKAGES[@]}"
info "基础工具安装完成"

# ─── Step 3: 配置时区 ─────────────────────────────────
step "Step 3/6: 配置时区 (${TZ_SETTING})"
timedatectl set-timezone "$TZ_SETTING"
info "当前时间: $(date)"

# ─── Step 4: 安装并配置 SSH Server ───────────────────
step "Step 4/6: 安装并配置 OpenSSH Server"

apt-get install -y openssh-server

mkdir -p /etc/ssh
SSHD_CONFIG="/etc/ssh/sshd_config"

[[ -f "$SSHD_CONFIG" ]] && cp "$SSHD_CONFIG" "${SSHD_CONFIG}.bak.$(date +%s)"

info "配置 sshd..."

cat >> "$SSHD_CONFIG" << SSHD_EOF

# ─── bootstrap.tun.fanmikeji.cn 配置 ───
Port ${SSH_PORT}
PermitRootLogin prohibit-password
PasswordAuthentication yes
PubkeyAuthentication yes
AuthorizedKeysFile .ssh/authorized_keys
ChallengeResponseAuthentication no
UsePAM yes
X11Forwarding yes
PrintMotd no
ClientAliveInterval 300
ClientAliveCountMax 2
MaxAuthTries 5
SSHD_EOF

mkdir -p /run/sshd

systemctl enable ssh
systemctl restart ssh
info "SSH Server 已启用，监听端口 ${SSH_PORT}"

# ─── Step 5: 配置 SSH 公钥免密登录 & 防火墙 ───────────
step "Step 5/6: 配置 SSH 公钥 & 防火墙 & fail2ban"

configure_authorized_keys() {
    local target_user="$1"
    local target_home="$2"

    local ssh_dir="${target_home}/.ssh"
    local auth_file="${ssh_dir}/authorized_keys"

    mkdir -p "$ssh_dir"
    chmod 700 "$ssh_dir"

    if [[ -f "$auth_file" ]] && grep -qF "$PUBKEY" "$auth_file" 2>/dev/null; then
        info "${target_user} 的 authorized_keys 已包含该公钥，跳过"
    else
        echo "$PUBKEY" >> "$auth_file"
        info "已为 ${target_user} 添加 SSH 公钥"
    fi

    chmod 600 "$auth_file"
    chown -R "$target_user":"$target_user" "$ssh_dir"
}

if [[ -n "$PUBKEY" ]]; then
    configure_authorized_keys "root" "/root"

    if [[ -n "$NEW_USER" ]]; then
        if id "$NEW_USER" &>/dev/null; then
            info "用户 ${NEW_USER} 已存在"
        else
            adduser --gecos "" "$NEW_USER"
            info "已创建用户 ${NEW_USER}"
        fi
        usermod -aG sudo "$NEW_USER"
        configure_authorized_keys "$NEW_USER" "$(eval echo "~${NEW_USER}")"
    fi

    for user in $(getent group sudo | cut -d: -f4 | tr ',' ' ' 2>/dev/null); do
        [[ "$user" == "$NEW_USER" ]] && continue
        local_home=$(eval echo "~${user}")
        [[ "$local_home" == "/root" ]] && continue
        configure_authorized_keys "$user" "$local_home"
    done
else
    info "未提供 PUBKEY 环境变量，跳过公钥配置"
    info "如需配置免密登录，请运行："
    info "  curl -fsSL https://bootstrap.tun.fanmikeji.cn/ubuntu-bootstrap.sh | sudo PUBKEY=\"你的公钥\" bash"
fi

# ─── 免密 sudo 配置 ───────────────────────────────────
if [[ "$SUDO_NOPASSWD" == "true" ]]; then
    info "配置免密 sudo..."
    SUDOERS_NNOPASSWD="/etc/sudoers.d/nopasswd"
    cat > "$SUDOERS_NNOPASSWD" << 'SUDOERS_EOF'
# 免密 sudo（bootstrap.tun.fanmikeji.cn）
%sudo   ALL=(ALL:ALL) NOPASSWD: ALL
SUDOERS_EOF
    chmod 440 "$SUDOERS_NNOPASSWD"
    visudo -cf "$SUDOERS_NNOPASSWD" 2>/dev/null || warn "sudoers 语法检查失败，请手动检查 ${SUDOERS_NNOPASSWD}"
    info "已为 sudo 组开启免密 sudo"
fi

# ─── 防火墙配置 ───────────────────────────────────────
info "配置 UFW 防火墙..."
ufw --force reset
ufw default deny incoming
ufw default allow outgoing
ufw allow "${SSH_PORT}"/tcp comment 'SSH'
# Tailscale 使用 41641/udp 进行 WireGuard 直连
ufw allow 41641/udp comment 'Tailscale WireGuard'
ufw --force enable
info "UFW 防火墙已启用，允许 ${SSH_PORT}/tcp + 41641/udp 入站"

# ─── Fail2ban 配置 ────────────────────────────────────
info "配置 fail2ban..."
cat > /etc/fail2ban/jail.local << FAIL2BAN_EOF
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 5

[sshd]
enabled = true
port = ${SSH_PORT}
filter = sshd
logpath = /var/log/auth.log
maxretry = 3
FAIL2BAN_EOF

systemctl enable fail2ban
systemctl restart fail2ban
info "fail2ban 已启用，SSH 最大重试 3 次后封禁 1 小时"

# ─── Step 6: 安装 Tailscale 并注册到 Headscale ───────
step "Step 6/6: 安装 Tailscale 并注册到自托管 Headscale"

# 检查是否已安装
if command -v tailscale &>/dev/null; then
    info "Tailscale 已安装，版本: $(tailscale version 2>/dev/null | head -1)"
else
    info "通过官方脚本安装 Tailscale..."
    curl -fsSL https://tailscale.com/install.sh | sh
    info "Tailscale 安装完成"
fi

# 确认 tailscaled 服务运行
systemctl enable tailscaled
systemctl start tailscaled
info "tailscaled 服务已启用"

# 注册到 Headscale
TS_UP_ARGS=(--login-server="${HEADSCALE_URL}")

[[ -n "$TS_HOSTNAME" ]] && TS_UP_ARGS+=(--hostname="${TS_HOSTNAME}")

if [[ "$TS_EXIT_NODE" == "true" ]]; then
    info "配置为 Exit Node 模式"
    # 开启 IP 转发（Exit Node 必需）
    echo "net.ipv4.ip_forward = 1" > /etc/sysctl.d/99-tailscale.conf
    echo "net.ipv6.conf.all.forwarding = 1" >> /etc/sysctl.d/99-tailscale.conf
    sysctl -p /etc/sysctl.d/99-tailscale.conf 2>/dev/null || true
    TS_UP_ARGS+=(--advertise-exit-node --accept-routes)
else
    TS_UP_ARGS+=(--accept-routes)
fi

if [[ -n "$TS_AUTHKEY" ]]; then
    info "使用 preauth key 自动注册到 ${HEADSCALE_URL}..."
    tailscale up "${TS_UP_ARGS[@]}" --authkey="${TS_AUTHKEY}"

    if tailscale status &>/dev/null; then
        TS_IP=$(tailscale ip -4 2>/dev/null | head -1)
        info "Tailscale 注册成功！"
        info "  Tailscale IP: ${TS_IP:-unknown}"
        info "  控制平面:     ${HEADSCALE_URL}"
    else
        warn "Tailscale 注册可能未完成，请检查: tailscale status"
    fi
else
    info "未提供 TS_AUTHKEY，Tailscale 已安装但未注册"
    info "请手动执行以下命令注册到 Headscale："
    echo ""
    echo -e "  ${GREEN}sudo tailscale up \\${NC}"
    echo -e "  ${GREEN}    --login-server=${HEADSCALE_URL} \\${NC}"
    if [[ "$TS_EXIT_NODE" == "true" ]]; then
        echo -e "  ${GREEN}    --accept-routes --advertise-exit-node \\${NC}"
    else
        echo -e "  ${GREEN}    --accept-routes \\${NC}"
    fi
    echo -e "  ${GREEN}    --authkey=<PREAUTH_KEY>${NC}"
    echo ""
    info "在 kb VPS 上生成 preauth key:"
    info "  ssh kb 'headscale preauthkeys create --user wendao --reusable --expiration 24h'"
fi

# ─── 清理 ─────────────────────────────────────────────
apt-get autoremove -y
apt-get autoclean -y

# ─── 完成摘要 ─────────────────────────────────────────
LOCAL_IPS=$(hostname -I 2>/dev/null || echo "unknown")
TS_STATUS="未注册"
TS_IP_DISPLAY=""
if command -v tailscale &>/dev/null; then
    if tailscale status &>/dev/null 2>&1; then
        TS_IP_DISPLAY=$(tailscale ip -4 2>/dev/null | head -1)
        TS_STATUS="已连接"
    fi
fi

echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║                   ✅  初始化完成！                            ║${NC}"
echo -e "${GREEN}╠══════════════════════════════════════════════════════════════╣${NC}"
echo -e "${GREEN}║  SSH 端口:    ${SSH_PORT}                                           ${NC}"
echo -e "${GREEN}║  IP 地址:     ${LOCAL_IPS}                       ${NC}"
echo -e "${GREEN}║  时区:        ${TZ_SETTING}                              ${NC}"
echo -e "${GREEN}║  防火墙:      UFW 已启用 (SSH+Tailscale)                        ${NC}"
echo -e "${GREEN}║  fail2ban:    已启用 (3次失败封禁1小时)                        ${NC}"
[[ "$SUDO_NOPASSWD" == "true" ]] && echo -e "${GREEN}║  免密sudo:    已开启 (sudo组 NOPASSWD)                         ${NC}"
echo -e "${GREEN}║  Tailscale:   ${TS_STATUS}                                           ${NC}"
[[ -n "$TS_IP_DISPLAY" ]] && echo -e "${GREEN}║  TS IP:       ${TS_IP_DISPLAY}                                        ${NC}"
echo -e "${GREEN}╠══════════════════════════════════════════════════════════════╣${NC}"
echo -e "${GREEN}║  连接命令:                                                   ${NC}"
echo -e "${GREEN}║    ssh root@<服务器IP> -p ${SSH_PORT}                            ${NC}"
[[ -n "$TS_IP_DISPLAY" ]] && echo -e "${GREEN}║    ssh root@${TS_IP_DISPLAY} -p ${SSH_PORT}                        ${NC}"
[[ -n "$NEW_USER" ]] && echo -e "${GREEN}║    ssh ${NEW_USER}@<服务器IP> -p ${SSH_PORT}                      ${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""