#!/bin/bash
# VPS 安全加固脚本 - Debian 12
# 用法: bash init.sh --user <username>

# set -eo pipefail

SSH_PORT=10010
KEY_TYPE="ed25519"

readonly R='\033[0;31m' G='\033[0;32m' Y='\033[1;33m' B='\033[0;34m' C='\033[0;36m' N='\033[0m'
info() { echo -e "${G}✓${N} $1"; }
warn() { echo -e "${Y}!${N} $1"; }
err() { echo -e "${R}✗${N} $1" >&2; exit 1; }

show_help() {
    cat << EOF
用法: bash init.sh --user <username>

选项:
    --user <username>    用户名（必需）
    -h, --help           帮助

示例: bash init.sh --user imhcg
EOF
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --user)
                NEW_USER="${2:-}"
                [[ -z "$NEW_USER" ]] && err "--user 需要值"
                shift 2
                ;;
            -h|--help) show_help; exit 0 ;;
            *) err "未知参数: $1" ;;
        esac
    done
    [[ -z "${NEW_USER:-}" ]] && err "缺少 --user"
}

check_env() {
    [[ $EUID -ne 0 ]] && err "请用 root 运行"
    local is_debian=false
    if [[ -f /etc/os-release ]]; then
        grep -qE "^ID=debian" /etc/os-release 2>/dev/null && is_debian=true
    fi
    if [[ "$is_debian" != "true" ]]; then
        warn "非 Debian 系统"
        read -p "继续? [y/N]: " c
        [[ "$c" =~ ^[Yy]$ ]] || exit 1
    fi
    return 0
}

set_host() {
    local cur=$(hostname)
    echo ""
    read -p "主机名 [${cur}]: " input </dev/tty
    HOSTNAME="${input:-$cur}"
    [[ "$HOSTNAME" != "$cur" ]] && hostnamectl set-hostname "$HOSTNAME" 2>/dev/null || true
    info "主机: $HOSTNAME"
}

ensure_iproute2() {
    if ! command -v ip &>/dev/null; then
        apt-get update -qq >/dev/null 2>&1
        apt-get install -y -qq iproute2 >/dev/null 2>&1
    fi
}

get_default_ip() {
    ensure_iproute2
    local ip
    ip=$(ip route get 1 2>/dev/null | awk '{print $7; exit}')
    [[ -z "$ip" ]] && ip=$(hostname -I 2>/dev/null | awk '{print $1}')
    [[ -z "$ip" ]] && ip=$(ifconfig 2>/dev/null | grep -Eo 'inet (addr:)?([0-9]*\.){3}[0-9]*' | grep -Eo '([0-9]*\.){3}[0-9]*' | grep -v '127.0.0.1' | head -1)
    echo "$ip"
}

set_ip() {
    local default_ip
    default_ip=$(get_default_ip)
    [[ -z "$default_ip" ]] && default_ip=""
    echo ""
    read -p "服务器 IP [${default_ip:-手动输入}]: " input </dev/tty
    SERVER_IP="${input:-$default_ip}"
    [[ -z "$SERVER_IP" ]] && err "需要服务器 IP"
    info "IP: $SERVER_IP"
}

backup() {
    [[ -f "$1" && ! -f "${1}.backup" ]] && cp "$1" "${1}.backup"
}

create_user() {
    # 确保必要工具已安装
    local tools="sudo openssl"
    for tool in $tools; do
        if ! command -v $tool &>/dev/null; then
            apt-get update -qq >/dev/null 2>&1
            apt-get install -y -qq $tool >/dev/null 2>&1
            info "$tool 已安装"
        fi
    done

    if id "$NEW_USER" &>/dev/null; then
        warn "用户 $NEW_USER 已存在，轮换密钥"
        read -p "继续? [y/N]: " c </dev/tty
        [[ "$c" =~ ^[Yy]$ ]] || err "已取消"
        groups "$NEW_USER" | grep -qw sudo || usermod -aG sudo "$NEW_USER" 2>/dev/null
        echo "$NEW_USER:$(openssl rand -base64 32)" | chpasswd 2>/dev/null
        info "用户已更新"
    else
        useradd -m -s /bin/bash "$NEW_USER" 2>/dev/null
        echo "$NEW_USER:$(openssl rand -base64 32)" | chpasswd 2>/dev/null
        usermod -aG sudo "$NEW_USER" 2>/dev/null
        info "用户已创建"
    fi

    # 配置 sudo 免密
    echo "$NEW_USER ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/$NEW_USER
    chmod 440 /etc/sudoers.d/$NEW_USER
    info "sudo 免密已配置"
}

gen_key() {
    local dir="/home/$NEW_USER/.ssh"
    local name="id_${NEW_USER}_${KEY_TYPE}"
    local file="$dir/$name"

    [[ -f "$file" ]] && {
        local t=$(date +%Y%m%d_%H%M%S)
        mv "$file" "${file}.old.$t" 2>/dev/null || true
        mv "${file}.pub" "${file}.pub.old.$t" 2>/dev/null || true
    }

    mkdir -p "$dir"
    chmod 700 "$dir"
    ssh-keygen -t "$KEY_TYPE" -a 100 -f "$file" -N "" -C "$NEW_USER@$HOSTNAME" 2>/dev/null
    cp "${file}.pub" "$dir/authorized_keys"
    chmod 600 "$dir/authorized_keys"
    chown -R "$NEW_USER:$NEW_USER" "$dir"
    info "密钥已生成"
}

harden_ssh() {
    backup "/etc/ssh/sshd_config"
    cat > /etc/ssh/sshd_config << EOF
# VPS 安全加固配置
Port $SSH_PORT
PermitRootLogin no
PasswordAuthentication no
PubkeyAuthentication yes
PermitEmptyPasswords no
MaxAuthTries 3
MaxSessions 2
X11Forwarding no
Ciphers chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com
MACs hmac-sha2-512-etm@openssh.com,hmac-sha2-256-etm@openssh.com
KexAlgorithms curve25519-sha256@libssh.org,diffie-hellman-group-exchange-sha256
ClientAliveInterval 300
ClientAliveCountMax 2
LoginGraceTime 60

# SFTP 支持
Subsystem sftp /usr/lib/openssh/sftp-server
EOF
    sshd -t 2>/dev/null || err "SSH 配置错误"
    systemctl restart sshd 2>/dev/null || service ssh restart 2>/dev/null
    info "SSH 已加固 (端口: $SSH_PORT)"
}

setup_ufw() {
    command -v ufw &>/dev/null || {
        apt-get update -qq >/dev/null 2>&1
        apt-get install -y -qq ufw >/dev/null 2>&1
    }
    ufw default deny incoming >/dev/null 2>&1
    ufw default allow outgoing >/dev/null 2>&1
    ufw allow "$SSH_PORT/tcp" comment 'SSH' >/dev/null 2>&1
    ufw --force enable >/dev/null 2>&1
    info "防火墙已启用"
}

setup_fail2ban() {
    command -v fail2ban-server &>/dev/null || apt-get install -y -qq fail2ban >/dev/null 2>&1
    cat > /etc/fail2ban/jail.local << EOF
[DEFAULT]
bantime = 1h
findtime = 10m
maxretry = 3

[sshd]
enabled = true
port = $SSH_PORT
filter = sshd
logpath = /var/log/auth.log
maxretry = 3
EOF
    systemctl enable fail2ban >/dev/null 2>&1
    systemctl restart fail2ban >/dev/null 2>&1
    info "fail2ban 已启用"
}

optimize_dns() {
    # 配置 DNS，使用 Cloudflare 和 Google 的 IPv4/IPv6 DNS
    cat > /etc/resolv.conf << 'EOF'
nameserver 1.1.1.1
nameserver 8.8.8.8
nameserver 2606:4700:4700::1111
nameserver 2001:4860:4860::8888
EOF
    # 防止被 NetworkManager 覆盖
    if [[ -f /etc/NetworkManager/NetworkManager.conf ]]; then
        grep -q "dns=none" /etc/NetworkManager/NetworkManager.conf || {
            sed -i '/\[main\]/a dns=none' /etc/NetworkManager/NetworkManager.conf
            systemctl restart NetworkManager 2>/dev/null || true
        }
    fi
    info "DNS 已优化"
}

check_swap() {
    # 检查是否已有 swap
    if swapon --show 2>/dev/null | grep -q "swap"; then
        info "swap 已存在"
        return 0
    fi
    
    # 检查内存大小
    local mem_total
    mem_total=$(free -m 2>/dev/null | awk '/^Mem:/{print $2}')
    [[ -z "$mem_total" ]] && mem_total=0
    
    # 内存小于 1G 时创建 swap
    if [[ "$mem_total" -lt 1024 ]]; then
        warn "内存 ${mem_total}M，尝试创建 1G swap"
        # 尝试创建 swap 文件，部分虚拟化环境可能不支持
        if fallocate -l 1G /swapfile 2>/dev/null || dd if=/dev/zero of=/swapfile bs=1M count=1024 2>/dev/null; then
            chmod 600 /swapfile
            mkswap /swapfile >/dev/null 2>&1
            swapon /swapfile 2>/dev/null && {
                echo "/swapfile none swap sw 0 0" >> /etc/fstab
                info "swap 已创建"
            } || {
                rm -f /swapfile
                warn "swap 创建失败，可能虚拟化不支持"
            }
        else
            warn "swap 创建失败，可能虚拟化不支持"
        fi
    else
        info "内存充足 (${mem_total}M)，跳过 swap"
    fi
}

limit_logs() {
    # 限制 journald 日志大小为 100M
    if [[ -f /etc/systemd/journald.conf ]]; then
        sed -i 's/#SystemMaxUse=/SystemMaxUse=/' /etc/systemd/journald.conf
        sed -i 's/SystemMaxUse=.*/SystemMaxUse=100M/' /etc/systemd/journald.conf
        grep -q "SystemMaxUse=" /etc/systemd/journald.conf || echo "SystemMaxUse=100M" >> /etc/systemd/journald.conf
        systemctl restart systemd-journald 2>/dev/null || true
    fi
    
    # 配置 logrotate 限制 /var/log
    cat > /etc/logrotate.d/vps-limit << 'EOF'
/var/log/*.log {
    rotate 3
    size 10M
    compress
    missingok
    notifempty
    delaycompress
}
EOF
    info "日志已限制"
}

show_config() {
    local key_name="id_${NEW_USER}_${KEY_TYPE}"
    local key_file="/home/$NEW_USER/.ssh/$key_name"
    local local_name="${HOSTNAME}"
    local b64=$(base64 -w 0 "$key_file")

    # SSH config 内容
    local ssh_config="Host $HOSTNAME
    HostName $SERVER_IP
    Port $SSH_PORT
    User $NEW_USER
    IdentityFile ~/.ssh/$local_name
    IdentitiesOnly yes"

    echo ""
    echo -e "${G}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${N}"
    echo -e "${G}  配置完成${N}"
    echo -e "${G}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${N}"
    echo "  主机: $HOSTNAME"
    echo "  IP:   $SERVER_IP"
    echo "  用户: $NEW_USER"
    echo "  端口: $SSH_PORT"
    
    echo ""
    echo -e "${Y}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${N}"
    echo -e "${Y}  配置 SSH (选一种方式)${N}"
    echo -e "${Y}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${N}"
    echo ""
    echo -e "${C}macOS/Linux 一键配置:${N}"
    echo -e "${Y}echo '$b64' | base64 -d > ~/.ssh/$local_name && chmod 600 ~/.ssh/$local_name && cat >> ~/.ssh/config << 'EOF'${N}"
    echo -e "${Y}$ssh_config${N}"
    echo -e "${Y}EOF${N}"
    echo ""
    echo -e "${C}Windows PowerShell 一键配置:${N}"
    echo -e "${Y}[System.IO.File]::WriteAllBytes(\"\$env:USERPROFILE/.ssh/$local_name\", [System.Convert]::FromBase64String(\"$b64\")); Add-Content -Path \$env:USERPROFILE/.ssh/config -Value @'${N}"
    echo -e "${Y}$ssh_config${N}"
    echo -e "${Y}'@${N}"
    echo ""
    
    echo -e "${C}手动配置:${N}"
    echo -e "${C}手动保存下面的密钥到 ~/.ssh/$local_name:${N}"
    echo ""
    cat "$key_file"
    echo ""
    echo -e "${C}手动添加 SSH 配置到 ~/.ssh/config:${N}"
    echo -e "${Y}Host $HOSTNAME${N}"
    echo -e "${Y}    HostName $SERVER_IP${N}"
    echo -e "${Y}    Port $SSH_PORT${N}"
    echo -e "${Y}    User $NEW_USER${N}"
    echo -e "${Y}    IdentityFile ~/.ssh/$local_name${N}"
    echo -e "${Y}    IdentitiesOnly yes${N}"
    echo ""
    echo -e "${C}一键连接:${N}"
    echo -e "${Y}ssh $HOSTNAME${N}"
    echo ""

    rm -f "$key_file"

    warn "务必保存私钥后再关闭终端"
}

main() {
    parse_args "$@"
    check_env || exit 1
    set_host || exit 1
    set_ip || exit 1

    echo ""
    echo -e "${B}开始加固...${N}"
    echo ""

    optimize_dns
    check_swap
    create_user
    gen_key
    harden_ssh
    setup_ufw
    setup_fail2ban
    limit_logs
    show_config
}

main "$@"