#!/bin/bash

# ============================================
# 资源消耗脚本（普通用户版）
# 用途：吃满 CPU 和内存
# 警告：可能导致系统卡死或 OOM
# ============================================

# 配置参数
CPU_CORES=$(nproc)                    # 自动获取 CPU 核心数
MEMORY_MB=2048                        # 要吃多少内存（MB），根据实际情况调整
SCRIPT_PID=$$

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

# 清理函数（退出时自动杀掉所有子进程）
cleanup() {
    echo -e "${YELLOW}[*] 正在清理所有子进程...${NC}"
    # 杀掉所有子进程
    pkill -P $$ 2>/dev/null
    # 杀掉所有 while 循环进程
    pkill -f "while :; do :; done" 2>/dev/null
    # 杀掉所有 dd 和 cat 进程
    pkill -f "dd if=/dev/zero" 2>/dev/null
    pkill -f "cat > /dev/null" 2>/dev/null
    echo -e "${GREEN}[✓] 清理完成${NC}"
    exit 0
}

# 捕获退出信号
trap cleanup SIGINT SIGTERM EXIT

# 显示 banner
echo -e "${RED}"
echo "╔═══════════════════════════════════════════╗"
echo "║       资源消耗脚本 (普通用户版)           ║"
echo "║   CPU: 将吃满全部 ${CPU_CORES} 个核心          ║"
echo "║   内存: 将吃掉 ${MEMORY_MB} MB                  ║"
echo "╚═══════════════════════════════════════════╝"
echo -e "${NC}"
echo -e "${YELLOW}[!] 警告：此操作可能导致系统卡死！${NC}"
echo -e "${YELLOW}[!] 按 Ctrl+C 可停止脚本${NC}"
echo ""
sleep 2

# ============================================
# 1. 吃满 CPU（开满所有核心）
# ============================================
echo -e "${GREEN}[*] 启动 CPU 消耗进程...${NC}"
for i in $(seq 1 $CPU_CORES); do
    # 每个核心跑一个无限循环
    while :; do :; done &
    echo "   启动 CPU 核心 $i 的进程 (PID: $!)"
done

# ============================================
# 2. 吃满内存
# ============================================
echo -e "${GREEN}[*] 启动内存消耗进程 (${MEMORY_MB} MB)...${NC}"

# 方法A：使用 tail + cat 吃内存（推荐，兼容性好）
tail -c ${MEMORY_MB}M /dev/zero 2>/dev/null | cat > /dev/null &
MEM_PID=$!
echo "   内存消耗进程 PID: $MEM_PID"

# 方法B：备选方案，用 dd（如果上面不行）
# dd if=/dev/zero of=/dev/null bs=1M count=$MEMORY_MB 2>/dev/null &
# MEM_PID=$!

# ============================================
# 3. 可选：吃磁盘 I/O（增加系统压力）
# ============================================
echo -e "${GREEN}[*] 是否需要额外消耗磁盘 I/O？(输入 y/n)${NC}"
read -t 5 answer 2>/dev/null || answer="n"
if [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
    echo -e "${GREEN}[*] 启动磁盘 I/O 消耗...${NC}"
    dd if=/dev/zero of=/tmp/hogfile bs=1M count=500 2>/dev/null &
    IO_PID=$!
    echo "   I/O 进程 PID: $IO_PID"
fi

# ============================================
# 4. 显示状态
# ============================================
echo ""
echo -e "${GREEN}[✓] 所有资源消耗进程已启动${NC}"
echo -e "${YELLOW}[*] 按 Ctrl+C 停止所有进程${NC}"
echo ""

# 循环显示当前资源使用情况
while true; do
    # CPU 负载（1分钟平均）
    LOAD=$(uptime | awk -F'load average:' '{print $2}' | cut -d, -f1 | xargs)
    # 内存使用（仅显示当前用户的进程内存占用估算）
    MEM_USAGE=$(ps -u $USER --no-headers -o rss 2>/dev/null | awk '{sum+=$1} END {printf "%.1f", sum/1024}')
    
    echo -ne "\r[$(date +%H:%M:%S)] CPU 负载: $LOAD | 用户 $USER 内存占用: ${MEM_USAGE:-0} MB    "
    sleep 2
done