HeXi/hexi/plugins/nonebot_plugin_web_hub/dashboard.py

# -*- coding: utf-8 -*-
"""统一 Web 管理台 Dashboard 数据采集。

数据来源与 nonebot_plugin_picstatus 类似：
- bot 账号 / WS 连通性：从 OneBot V11 适配器的 bots / connections 读取
- host 性能：psutil（CPU/内存/磁盘/进程）+ platform + nonebot 版本
仅在 /hub/api/dashboard 被调用时执行，不做常驻采样。
"""

from __future__ import annotations

import asyncio
import platform
import time

import psutil
from nonebot import get_bots, logger

try:
    from cpuinfo import get_cpu_info
except Exception:  # pragma: no cover
    get_cpu_info = None

_MODULE_START = time.time()


async def _collect_cpu() -> dict:
    try:
        per_core = await asyncio.to_thread(psutil.cpu_percent, 0.3, True)
    except TypeError:
        # 兼容部分平台不接受 percpu 参数
        per_core = []
    if per_core:
        percent = sum(per_core) / len(per_core)
    else:
        percent = await asyncio.to_thread(psutil.cpu_percent, 0.3)
    info: dict = {
        "percent": round(float(percent), 1),
        "per_core": [round(float(x), 1) for x in (per_core or [])],
        "logical": psutil.cpu_count(),
        "physical": psutil.cpu_count(logical=False),
        "brand": "",
    }
    if get_cpu_info:
        try:
            brand = get_cpu_info().get("brand_raw", "")
            info["brand"] = brand.split("@", maxsplit=1)[0].strip()
        except Exception:
            info["brand"] = ""
    try:
        info["load_avg"] = [round(float(x), 2) for x in psutil.getloadavg()]
    except Exception:
        info["load_avg"] = None
    return info


async def _collect_memory() -> dict:
    mem = psutil.virtual_memory()
    try:
        proc = psutil.Process()
        rss = proc.memory_info().rss
        proc_mem_percent = proc.memory_percent()
    except Exception:
        rss = 0
        proc_mem_percent = 0.0
    return {
        "percent": round(float(mem.percent), 1),
        "used": int(mem.used),
        "total": int(mem.total),
        "available": int(mem.available),
        "process_rss": int(rss),
        "process_percent": round(float(proc_mem_percent), 1),
    }


async def _collect_process() -> dict:
    try:
        proc = psutil.Process()
        cpu = await asyncio.to_thread(proc.cpu_percent, 0.2)
        info = {
            "pid": proc.pid,
            "name": proc.name(),
            "rss": int(proc.memory_info().rss),
            "cpu_percent": round(float(cpu), 1),
            "memory_percent": round(float(proc.memory_percent()), 1),
            "create_time": int(proc.create_time()),
        }
    except Exception as e:
        logger.warning(f"采集进程信息失败: {type(e).__name__}: {e}")
        info = {}
    return info


async def _collect_system() -> dict:
    boot = psutil.boot_time()
    now = time.time()
    return {
        "hostname": platform.node(),
        "os": platform.platform(),
        "python": f"{platform.python_implementation()} {platform.python_version()}",
        "nonebot": _nonebot_version(),
        "boot_time": int(boot),
        "system_uptime_seconds": int(now - boot),
        "bot_uptime_seconds": int(now - _MODULE_START),
    }


def _nonebot_version() -> str:
    try:
        import nonebot

        return nonebot.__version__
    except Exception:
        return ""


async def _collect_disk() -> list:
    parts = []
    for part in psutil.disk_partitions():
        try:
            usage = psutil.disk_usage(part.mountpoint)
        except Exception:
            continue
        parts.append(
            {
                "mount": part.mountpoint,
                "device": part.device,
                "percent": round(float(usage.percent), 1),
                "used": int(usage.used),
                "total": int(usage.total),
            }
        )
    return parts


async def _collect_bots() -> list:
    try:
        bots = get_bots().values()
    except Exception as e:  # noqa: BLE001
        logger.warning(f"获取 bot 列表失败: {type(e).__name__}: {e}")
        return []
    result = []
    for bot in bots:
        adapter = getattr(bot, "adapter", None)
        connections = getattr(adapter, "connections", {}) or {}
        ws_connected = bool(connections.get(bot.self_id))
        item = {
            "self_id": bot.self_id,
            "adapter": adapter.get_name() if adapter else "",
            "nick": bot.self_id,
            "avatar": f"https://q1.qlogo.cn/g?b=qq&nk={bot.self_id}&s=640",
            "ws_connected": ws_connected,
            "online": None,
            "msg_recv": None,
            "msg_sent": None,
        }
        if ws_connected:
            try:
                login = await bot.get_login_info()
                item["nick"] = login.get("nickname") or item["nick"]
            except Exception as e:  # noqa: BLE001
                logger.warning(
                    f"获取登录信息失败({bot.self_id}): {type(e).__name__}: {e}"
                )
            try:
                status = await bot.get_status()
                item["online"] = status.get("online")
                stat = status.get("stat") or {}
                item["msg_recv"] = stat.get("message_received") or stat.get(
                    "MessageReceived"
                )
                item["msg_sent"] = stat.get("message_sent") or stat.get("MessageSent")
            except Exception as e:  # noqa: BLE001
                logger.warning(
                    f"获取连接状态失败({bot.self_id}): {type(e).__name__}: {e}"
                )
        result.append(item)
    return result


async def collect_dashboard() -> dict:
    """聚合采集，失败字段降级，不阻塞整体返回。"""
    cpu, memory, process, system, disk, bots = await asyncio.gather(
        _collect_cpu(),
        _collect_memory(),
        _collect_process(),
        _collect_system(),
        _collect_disk(),
        _collect_bots(),
    )
    return {
        "bots": bots,
        "system": system,
        "cpu": cpu,
        "memory": memory,
        "process": process,
        "disk": disk,
    }