HeXi/hexi/plugins/nonebot_plugin_picstatus/misc_statistics.py

from collections.abc import Callable
from datetime import datetime
from typing import Any

from cookit.loguru import log_exception_warning
from nonebot import get_driver, logger
from nonebot.adapters import Bot as BaseBot, Event as BaseEvent
from nonebot.message import event_preprocessor
from nonebot.typing import T_State
from nonebot_plugin_alconna import image_fetch
from nonebot_plugin_alconna.uniseg import Image
from nonebot_plugin_uninfo import User, get_interface

from .config import config

nonebot_run_time: datetime = datetime.now().astimezone()
bot_connect_time: dict[str, datetime] = {}
recv_num: dict[str, int] = {}
send_num: dict[str, int] = {}

bot_info_cache: dict[str, User] = {}
bot_avatar_cache: dict[str, bytes | None] = {}

driver = get_driver()

SEND_APIS: dict[str, list[str] | Callable[[str], bool]] = {
    # "BilibiliLive": [],  # 狗东西发消息不走 call_api
    "Console": ["send_msg"],
    "Ding": ["send"],
    "Discord": ["create_message"],
    "Feishu": ["im/v1/messages"],
    "Kaiheila": ["message_create", "directMessage_create"],
    "Minecraft": ["send_msg"],
    "mirai2": ["send_friend_message", "send_group_message", "send_temp_message"],
    "ntchat": lambda x: x.startswith("send_"),
    "OneBot V11": ["send_private_msg", "send_group_msg", "send_msg"],
    "OneBot V12": ["send_message"],
    "QQ": [
        "post_dms_messages",
        "post_messages",
        "post_c2c_messages",
        "post_c2c_files",
        "post_group_messages",
        "post_group_files",
    ],
    "RedProtocol": ["send_message", "send_fake_forward"],
    "Satori": ["message_create"],
    "Telegram": lambda x: x.startswith("send_"),
    "大别野": ["send_message"],
}


def method_is_send_msg(platform: str, name: str) -> bool:
    return (platform in SEND_APIS) and (
        (name in it) if isinstance((it := SEND_APIS[platform]), list) else it(name)
    )


if config.ps_count_message_sent_event:

    @event_preprocessor
    async def _(bot: BaseBot, event: BaseEvent):
        if (
            config.ps_count_message_sent_event
            and (
                (config.ps_count_message_sent_event is True)
                or bot.adapter.get_name() in config.ps_count_message_sent_event
            )
            and (
                (event.get_type() == "message_sent")
                or (
                    event.get_type() == "message" and event.get_user_id() == bot.self_id
                )
            )
        ):
            # logger.debug(f"Bot {bot.self_id} sent counter +1")
            send_num[bot.self_id] += 1


if config.ps_count_message_sent_event is not True:

    @BaseBot.on_called_api
    async def called_api(
        bot: BaseBot,
        exc: Exception | None,
        api: str,
        _: dict[str, Any],
        __: Any,
    ):
        if (
            (not exc)
            and (config.ps_count_message_sent_event is not True)
            and (
                (config.ps_count_message_sent_event is False)
                or (bot.adapter.get_name() not in config.ps_count_message_sent_event)
            )
            and method_is_send_msg(bot.adapter.get_name(), api)
        ):
            # logger.debug(f"Bot {bot.self_id} sent counter +1")
            send_num[bot.self_id] += 1


async def cache_bot_avatar(avatar: str, bot: BaseBot, event: BaseEvent, state: T_State):
    try:
        img = await image_fetch(event, bot, state, Image(url=avatar))
    except Exception as e:
        log_exception_warning(e, f"Failed to get avatar of bot {bot.self_id}")
        return None
    else:
        if not img:
            logger.warning(
                f"Cannot get avatar of bot {bot.self_id}"
                f" because image_fetch returned None",
            )
    bot_avatar_cache[bot.self_id] = img
    return img


async def cache_bot_info(bot: BaseBot):
    try:
        it = get_interface(bot)
        info = (await it.get_user(bot.self_id)) if it else None
    except Exception as e:
        logger.warning(f"Error when getting bot info: {e.__class__.__name__}: {e}")
    else:
        if info:
            bot_info_cache[bot.self_id] = info
            return info
    return None


@driver.on_bot_connect
async def _(bot: BaseBot):
    bot_connect_time[bot.self_id] = datetime.now().astimezone()
    if bot.self_id not in recv_num:
        recv_num[bot.self_id] = 0
    if (bot.self_id not in send_num) and (bot.adapter.get_name() in SEND_APIS):
        send_num[bot.self_id] = 0
    await cache_bot_info(bot)


@driver.on_bot_disconnect
async def _(bot: BaseBot):
    bot_connect_time.pop(bot.self_id, None)
    if config.ps_disconnect_reset_counter:
        recv_num.pop(bot.self_id, None)
        send_num.pop(bot.self_id, None)


@event_preprocessor
async def _(bot: BaseBot, event: BaseEvent):
    if event.get_type() == "message":
        recv_num[bot.self_id] += 1