LinuxAwesomeBot/monitor.py

"""Linux 系统监控：CPU、内存、温度、磁盘、网络、进程。"""
import os
import time
from pathlib import Path

import psutil


# ---------- 指标 ----------

def get_cpu() -> dict:
    load1, load5, load15 = _safe_getloadavg()
    return {
        "percent": psutil.cpu_percent(interval=None),
        "per_core": psutil.cpu_percent(interval=None, percpu=True),
        "load": [load1, load5, load15],
        "cores": psutil.cpu_count(logical=True),
    }


def _safe_getloadavg():
    try:
        return psutil.getloadavg()
    except (AttributeError, OSError):  # Windows
        return (0.0, 0.0, 0.0)


def get_memory() -> dict:
    vm = psutil.virtual_memory()
    sw = psutil.swap_memory()
    return {
        "used_gb": round(vm.used / 1024**3, 2),
        "total_gb": round(vm.total / 1024**3, 2),
        "percent": vm.percent,
        "swap_percent": sw.percent,
    }


def _wmi_temps() -> dict:
    """Windows: 通过 WMI 热区读取温度（部分主板支持，可能需要管理员权限）。"""
    try:
        import wmi  # 仅 Windows 安装
        zones = wmi.WMI(namespace="root\\wmi").MSAcpi_ThermalZoneTemperature()
        values = [z.CurrentTemperature for z in zones]
    except Exception:  # 类不存在 / 权限不足 / 非 Windows
        return {}
    # WMI 单位是 0.1 开尔文
    return {"ThermalZone": [
        {"label": f"zone{i}", "current": round(v / 10 - 273.15, 1)}
        for i, v in enumerate(values) if v
    ]}


def get_temps() -> dict:
    try:
        sensors = psutil.sensors_temperatures()
    except (AttributeError, OSError):
        sensors = {}
    out = {}
    for name, entries in sensors.items():
        out[name] = [
            {"label": e.label or f"temp{i}", "current": e.current}
            for i, e in enumerate(entries)
            if e.current is not None
        ]
    if not out and os.name == "nt":
        out = _wmi_temps()
    return out


def get_disk() -> dict:
    out = []
    for part in psutil.disk_partitions():
        if "loop" in part.device or not part.mountpoint:
            continue
        try:
            u = psutil.disk_usage(part.mountpoint)
            out.append({
                "device": part.device,
                "mount": part.mountpoint,
                "used_gb": round(u.used / 1024**3, 1),
                "total_gb": round(u.total / 1024**3, 1),
                "percent": u.percent,
            })
        except (PermissionError, OSError):
            continue
    return out


_net_prev = {"ts": None, "io": None}


def get_net(interval: float = 0.0) -> dict:
    """网络速率；interval>0 时阻塞采样该秒数，否则用上次调用差分。"""
    global _net_prev
    if interval > 0:
        io0 = psutil.net_io_counters()
        time.sleep(interval)
        io1 = psutil.net_io_counters()
        dt = interval
    else:
        now = time.monotonic()
        io1 = psutil.net_io_counters()
        if _net_prev["ts"] is None:
            _net_prev = {"ts": now, "io": io1}
            return {"up_mbps": None, "down_mbps": None}
        dt = now - _net_prev["ts"]
        io0 = _net_prev["io"]
        _net_prev = {"ts": now, "io": io1}
    if dt <= 0:
        return {"up_mbps": None, "down_mbps": None}
    return {
        "up_mbps": round((io1.bytes_sent - io0.bytes_sent) * 8 / dt / 1e6, 2),
        "down_mbps": round((io1.bytes_recv - io0.bytes_recv) * 8 / dt / 1e6, 2),
    }


_primed = False
_prime_ts = 0.0


def get_top_processes(n: int = 10) -> list:
    """按 CPU 占用排序的 TOP N 进程。首次（或距上次超过60秒）需先基准采样。"""
    global _primed, _prime_ts
    if not _primed or time.monotonic() - _prime_ts > 60:
        for p in psutil.process_iter(["pid"]):
            try:
                p.cpu_percent(interval=None)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        time.sleep(0.5)
        _primed = True
        _prime_ts = time.monotonic()
    procs = []
    for p in psutil.process_iter(["pid", "name"]):
        try:
            if p.info["pid"] in (0, 4):  # Windows 的 Idle/System 不计入排行
                continue
            cpu = p.cpu_percent(interval=None)
            mem = p.memory_info().rss / 1024**2
            procs.append({
                "pid": p.info["pid"],
                "name": (p.info["name"] or "?")[:24],
                "cpu": round(cpu, 1),
                "mem_mb": round(mem, 1),
            })
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            continue
    _prime_ts = time.monotonic()
    procs.sort(key=lambda x: x["cpu"], reverse=True)
    return procs[:n]


def get_top_memory(n: int = 10) -> list:
    """按内存占用排序的 TOP N 进程（无需 CPU 基准采样）。"""
    procs = []
    for p in psutil.process_iter(["pid", "name"]):
        try:
            if p.info["pid"] in (0, 4):
                continue
            procs.append({
                "pid": p.info["pid"],
                "name": (p.info["name"] or "?")[:24],
                "cpu": 0.0,
                "mem_mb": round(p.memory_info().rss / 1024**2, 1),
            })
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            continue
    procs.sort(key=lambda x: x["mem_mb"], reverse=True)
    return procs[:n]


def get_uptime() -> float:
    return time.time() - psutil.boot_time()


def get_status(top_n: int | None = None) -> dict:
    """系统总览。top_n 为 None 时不枚举进程（Windows 上较慢）。"""
    data = {
        "cpu": get_cpu(),
        "memory": get_memory(),
        "temps": get_temps(),
        "disk": get_disk(),
        "net": get_net(),
        "top": [],
        "uptime_s": get_uptime(),
    }
    if top_n is not None:
        data["top"] = get_top_processes(top_n)
    return data