#!/bin/bash
# Debian 按月流量超限自动关机脚本安装器
# 功能：每月1号自动清零、70GB预警、85GB关机、30秒检查间隔

set -e

echo "========================================"
echo "  按月流量超限自动关机脚本安装器"
echo "  预警阈值：70GB（发送一次全局警告）"
echo "  关机阈值：85GB（60秒后自动关机）"
echo "  检查间隔：30秒"
echo "  自动清零：每月1号0点"
echo "========================================"


echo -e "\n[0/4] 清理旧版本进程..."
pkill -f traffic_shutdown.sh 2>/dev/null || true


echo -e "\n[1/4] 创建按月流量监控脚本..."
cat > /usr/local/bin/traffic_shutdown.sh << 'EOF'
#!/bin/bash
# 配置参数
WARNING_GB=70                    # 预警阈值(GB)
SHUTDOWN_GB=85                   # 关机阈值(GB)
CHECK_INTERVAL=30                # 检查间隔(秒)
DATA_FILE="/var/lib/traffic_monitor/traffic_data"  # 数据存储文件

# 创建数据目录
mkdir -p /var/lib/traffic_monitor

# 初始化数据文件格式：累计流量(字节) 统计月份(YYYYMM) 预警已发送(0/1)
init_data_file() {
    local current_month=$(date +%Y%m)
    echo "0 $current_month 0" > "$DATA_FILE"
    echo "初始化数据文件，当前统计月份：$current_month"
}

# 获取当前所有网络接口的总流量(字节)
get_total_traffic() {
    local rx_bytes=$(grep -v 'lo:' /proc/net/dev | awk '{sum += $2} END {print sum}')
    local tx_bytes=$(grep -v 'lo:' /proc/net/dev | awk '{sum += $10} END {print sum}')
    echo $((rx_bytes + tx_bytes))
}

# 主程序开始
echo "按月流量监控服务已启动"
echo "预警阈值：${WARNING_GB}GB"
echo "关机阈值：${SHUTDOWN_GB}GB"
echo "检查间隔：${CHECK_INTERVAL}秒"
echo "数据文件：$DATA_FILE"

# 检查并初始化数据文件
if [ ! -f "$DATA_FILE" ]; then
    init_data_file
fi

# 读取数据文件
read total_bytes saved_month warning_sent < "$DATA_FILE"

# 检查月份是否变化，变化则自动清零
current_month=$(date +%Y%m)
if [ "$current_month" != "$saved_month" ]; then
    echo "检测到新月份：$current_month，自动清零流量统计"
    wall "⚠️  新月份开始，服务器流量统计已自动清零"
    init_data_file
    read total_bytes saved_month warning_sent < "$DATA_FILE"
fi

# 计算系统启动时的基准流量
current_system=$(get_total_traffic)
base_traffic=$((current_system - total_bytes))

# 主循环
while true; do
    # 再次检查月份（处理跨月运行情况）
    current_month=$(date +%Y%m)
    if [ "$current_month" != "$saved_month" ]; then
        echo "检测到新月份：$current_month，自动清零流量统计"
        wall "⚠️  新月份开始，服务器流量统计已自动清零"
        init_data_file
        read total_bytes saved_month warning_sent < "$DATA_FILE"
        current_system=$(get_total_traffic)
        base_traffic=$((current_system - total_bytes))
    fi

    # 获取当前流量并计算累计
    current=$(get_total_traffic)
    total_bytes=$((current - base_traffic))
    total_gb=$(echo "scale=2; $total_bytes / 1024 / 1024 / 1024" | bc)
    
    # 保存数据到文件
    echo "$total_bytes $saved_month $warning_sent" > "$DATA_FILE"
    
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] 当前累计流量：${total_gb}GB / ${SHUTDOWN_GB}GB (月份：$saved_month)"
    
    # 检查是否达到预警阈值（只发送一次）
    if (( $(echo "$total_gb >= $WARNING_GB" | bc -l) )) && [ "$warning_sent" -eq 0 ]; then
        echo "⚠️  流量已达到${WARNING_GB}GB预警阈值！"
        wall "⚠️  服务器流量预警：本月已使用${total_gb}GB，接近${SHUTDOWN_GB}GB限制，超过将自动关机！"
        warning_sent=1
        echo "$total_bytes $saved_month $warning_sent" > "$DATA_FILE"
    fi
    
    # 检查是否达到关机阈值
    if (( $(echo "$total_gb >= $SHUTDOWN_GB" | bc -l) )); then
        echo "⚠️  流量已超过${SHUTDOWN_GB}GB，系统将在60秒后关机！"
        wall "⚠️  紧急通知：服务器本月流量已超过${SHUTDOWN_GB}GB限制，系统将在60秒后自动关机！"
        sleep 60
        shutdown -h now
        exit 0
    fi
    
    sleep $CHECK_INTERVAL
done
EOF

# 2. 赋予执行权限
echo -e "\n[2/4] 赋予脚本执行权限..."
chmod +x /usr/local/bin/traffic_shutdown.sh

# 3. 后台启动脚本（nohup守护进程，日志重定向）
echo -e "\n[3/4] 启动后台监控服务..."
nohup /usr/local/bin/traffic_shutdown.sh > /var/log/traffic_monitor.log 2>&1 &
sleep 2

# 4. 配置开机自启（rc.local通用方案）
echo -e "\n[4/4] 配置开机自动启动..."
# 确保rc.local存在并可执行
if [ ! -f /etc/rc.local ]; then
    cat > /etc/rc.local << 'EOF'
#!/bin/sh -e
# 开机自动启动流量监控脚本
nohup /usr/local/bin/traffic_shutdown.sh > /var/log/traffic_monitor.log 2>&1 &
exit 0
EOF
    chmod +x /etc/rc.local
else
    # 若已存在，追加启动命令（避免重复）
    if ! grep -q "traffic_shutdown.sh" /etc/rc.local; then
        sed -i '/exit 0/d' /etc/rc.local
        echo 'nohup /usr/local/bin/traffic_shutdown.sh > /var/log/traffic_monitor.log 2>&1 &' >> /etc/rc.local
        echo 'exit 0' >> /etc/rc.local
    fi
fi

echo -e "\n========================================"
echo "✅ 安装/更新完成！"
echo ""
echo "📊 功能说明："
echo "• 每月1号0点自动清零流量统计"
echo "• 达到70GB时发送一次全局预警通知"
echo "• 达到85GB时60秒后自动关机"
echo "• 每30秒检查一次流量"
echo "• 服务器重启后流量统计不中断"
echo ""
echo "🔧 常用命令："
echo "查看实时日志：tail -f /var/log/traffic_monitor.log"
echo "查看本月流量：cat /var/lib/traffic_monitor/traffic_data | awk '{printf \"%.2f GB\n\", \$1 / 1024 / 1024 / 1024}'"
echo "查看统计月份：cat /var/lib/traffic_monitor/traffic_data | awk '{print \$2}'"
echo "重启监控服务：pkill -f traffic_shutdown.sh && nohup /usr/local/bin/traffic_shutdown.sh > /var/log/traffic_monitor.log 2>&1 &"
echo "手动清零流量：rm -f /var/lib/traffic_monitor/traffic_data && pkill -f traffic_shutdown.sh && nohup /usr/local/bin/traffic_shutdown.sh > /var/log/traffic_monitor.log 2>&1 &"
echo "修改阈值：编辑 /usr/local/bin/traffic_shutdown.sh 中的 WARNING_GB 和 SHUTDOWN_GB"
echo "========================================"