tg-spam-guard/handlers/spam.py

"""垃圾消息检测 handler"""
import asyncio
import logging
from datetime import datetime

from telegram import Update
from telegram.ext import ContextTypes, MessageHandler, filters

import config
from db.models import Message, Ban, Action, Whitelist
from services.spam_detector import detect
from web.api import broadcast_sse, json_safe

logger = logging.getLogger("spam_guard")

# 简单的防刷屏缓存：user_id -> [(timestamp, msg_id)]
_flood_cache: dict[int, list] = {}


async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """处理群消息：防刷屏 + 垃圾检测"""
    msg = update.effective_message
    if not msg or not msg.text:
        return

    user = update.effective_user
    if not user or user.is_bot:
        return

    username = user.username or ""
    first_name = user.first_name or ""
    text = msg.text
    logger.info(f"💬 聊天消息 [{username or first_name or user.id}]: {text[:120]}")

    # === 防刷屏 ===
    now = datetime.now().timestamp()
    if user.id not in _flood_cache:
        _flood_cache[user.id] = []
    _flood_cache[user.id].append((now, msg.message_id))
    # 清理过期
    _flood_cache[user.id] = [
        (t, mid) for t, mid in _flood_cache[user.id]
        if now - t < config.RATE_LIMIT_WINDOW
    ]
    if len(_flood_cache[user.id]) > config.RATE_LIMIT_MESSAGES:
        try:
            await msg.delete()
        except Exception:
            pass
        logger.warning(f"⚡ 防刷屏: @{username} ({user.id}) {len(_flood_cache[user.id])}条/{config.RATE_LIMIT_WINDOW}s")
        return

    # 记录消息并推送到 Web 管理页
    message_record = await Message.create(
        msg_id=msg.message_id,
        chat_id=update.effective_chat.id,
        user_id=user.id,
        username=username,
        first_name=first_name,
        text=text[:500],
    )
    await broadcast_sse("chat", json_safe({
        "id": message_record.id,
        "msg_id": msg.message_id,
        "chat_id": update.effective_chat.id,
        "user_id": user.id,
        "username": username,
        "first_name": first_name,
        "text": text[:500],
        "time": message_record.created_at,
        "spam": False,
    }))

    # === 垃圾检测 ===
    is_spam, reason = await detect(text, username, first_name)

    if is_spam:
        logger.warning(f"🚨 垃圾消息 [{username}]: {reason}")
        await message_record.update_from_dict({"is_spam": True, "spam_reason": reason}).save()
        await broadcast_sse("spam", json_safe({
            "id": message_record.id,
            "msg_id": msg.message_id,
            "chat_id": update.effective_chat.id,
            "user_id": user.id,
            "username": username,
            "first_name": first_name,
            "text": text[:500],
            "time": message_record.created_at,
            "spam": True,
            "reason": reason,
        }))

        # 删消息，并把记录标成已删除，避免前端继续显示“删除/删封”按钮
        try:
            await msg.delete()
            message_record.manually_deleted = True
            message_record.deleted_at = datetime.now()
            message_record.delete_reason = "auto_spam_delete"
            await message_record.save()
            await Action.create(action="DELETE", chat_id=update.effective_chat.id, user_id=user.id, username=username, details={"reason": reason, "text": text[:120]})
            await broadcast_sse("chat_delete", json_safe({
                "id": message_record.id,
                "msg_id": msg.message_id,
                "chat_id": update.effective_chat.id,
                "deleted_at": message_record.deleted_at,
            }))
        except Exception as e:
            logger.error(f"❌ 删除消息失败: {e}")

        # 管理员/白名单：只删消息，不封禁，避免误封可信用户
        if user.id in config.ADMIN_USER_IDS:
            await Action.create(action="SPAM_ADMIN_TEST", chat_id=update.effective_chat.id, user_id=user.id, username=username, details={"reason": reason})
            logger.warning(f"🧪 管理员垃圾规则测试命中，只删不封: [{username}] {reason}")
            return
        whitelist = await Whitelist.get_or_none(chat_id=update.effective_chat.id, user_id=user.id, enabled=True)
        if whitelist:
            await Action.create(action="SPAM_WHITELIST_DELETE_ONLY", chat_id=update.effective_chat.id, user_id=user.id, username=username, details={"reason": reason, "whitelist_id": whitelist.id})
            logger.warning(f"🛡️ 白名单用户命中垃圾规则，只删不封: [{username}] {reason}")
            return

        # 封禁
        try:
            await update.effective_chat.ban_member(user_id=user.id)
            unban_at = None

            ban_record = {
                "chat_id": update.effective_chat.id,
                "user_id": user.id,
                "username": username,
                "first_name": first_name,
                "reason": reason,
                "method": "auto",
                "duration_minutes": 0,
                "auto_unban": False,
                "unban_at": unban_at,
            }
            ban = await Ban.create(**ban_record)
            await Action.create(action="BAN", chat_id=update.effective_chat.id, user_id=user.id, username=username, details={"reason": reason})
            await broadcast_sse("ban", json_safe({
                "id": ban.id,
                "user_id": user.id,
                "username": username,
                "first_name": first_name,
                "reason": reason,
                "method": "auto",
                "duration_minutes": 0,
                "auto_unban": False,
                "unban_at": unban_at,
                "time": ban.created_at,
            }))
            logger.info(f"🔨 已永久封禁 [{username}]")
        except Exception as e:
            await Action.create(
                action="BAN_FAILED",
                chat_id=update.effective_chat.id,
                user_id=user.id,
                username=username,
                details={"reason": reason, "error": f"{type(e).__name__}: {e}"},
            )
            logger.error(f"❌ 封禁失败 [{username or user.id}]: {type(e).__name__}: {e}")

        # 标记为垃圾
        await Message.filter(chat_id=update.effective_chat.id, msg_id=msg.message_id).update(is_spam=True, spam_reason=reason)


def register_spam(app):
    app.add_handler(MessageHandler(filters.TEXT & filters.ChatType.GROUPS & ~filters.COMMAND, handle_message))