sshclient/core/monitor.py

"""
远程主机系统监控模块
通过 SSH 一次性采集 CPU/内存/磁盘/网络/负载指标。
Linux 用 /proc 和常用命令；macOS/BSD 走兼容路径。
"""
import re
import time
from typing import Optional

from .ssh_client import SSHConnection


class SystemMonitor:
    """远程主机的资源监控器（数据全部从 SSH 通道采集，不依赖 agent）"""

    # 一次性获取所有指标的脚本（Linux）
    _LINUX_METRICS_SCRIPT = r"""
echo "===CPU==="
# 第一次采样 1 秒间隔，用来计算差值
read cpu_user cpu_nice cpu_system cpu_idle cpu_iowait cpu_irq cpu_softirq cpu_steal < /proc/stat
sleep 1
read cpu_user2 cpu_nice2 cpu_system2 cpu_idle2 cpu_iowait2 cpu_irq2 cpu_softirq2 cpu_steal2 < /proc/stat
total1=$((cpu_user+cpu_nice+cpu_system+cpu_idle+cpu_iowait+cpu_irq+cpu_softirq+cpu_steal))
total2=$((cpu_user2+cpu_nice2+cpu_system2+cpu_idle2+cpu_iowait2+cpu_irq2+cpu_softirq2+cpu_steal2))
idle1=$cpu_idle; idle2=$cpu_idle2
dt=$((total2-total1)); di=$((idle2-idle1))
if [ $dt -gt 0 ]; then usage=$(( (1000*(dt-di)/dt+5)/10 )); else usage=0; fi
echo "CPU_USAGE=$usage"
echo "CPU_CORES=$(nproc 2>/dev/null || echo 1)"
echo "LOAD=$(cat /proc/loadavg | awk '{print $1,$2,$3}')"
echo "UPTIME=$(awk '{printf "%.0f",$1}' /proc/uptime)"
echo "===MEM==="
mem_total=$(awk '/MemTotal/{print $2}' /proc/meminfo)
mem_avail=$(awk '/MemAvailable/{print $2}' /proc/meminfo)
swap_total=$(awk '/SwapTotal/{print $2}' /proc/meminfo)
swap_free=$(awk '/SwapFree/{print $2}' /proc/meminfo)
if [ -z "$mem_avail" ]; then mem_avail=$((mem_total - $(awk '/^(Buffers|Cached|SReclaimable):/{s+=$2} END{print s}' /proc/meminfo))); fi
used=$((mem_total - mem_avail))
echo "MEM_TOTAL=$mem_total"
echo "MEM_USED=$used"
echo "MEM_AVAIL=$mem_avail"
echo "SWAP_TOTAL=$swap_total"
echo "SWAP_USED=$((swap_total-swap_free))"
echo "===DISK==="
df -PB1 -x tmpfs -x devtmpfs 2>/dev/null | awk 'NR>1 {printf "DISK|%s|%d|%d|%s\n",$NF,$2,$3,$5}'
echo "===NET==="
for iface in $(ls /sys/class/net/ 2>/dev/null | grep -v lo); do
  rx=$(cat /sys/class/net/$iface/statistics/rx_bytes 2>/dev/null || echo 0)
  tx=$(cat /sys/class/net/$iface/statistics/tx_bytes 2>/dev/null || echo 0)
  echo "NET|$iface|$rx|$tx"
done
echo "===PROC==="
# 进程统计
ps -e -o stat= 2>/dev/null | awk '{r+=($1~/^R/); s+=($1~/^S/); d+=($1~/^D/); z+=($1~/^Z/); t++} END{printf "PROC_TOTAL=%d\nPROC_RUNNING=%d\nPROC_SLEEP=%d\nPROC_DISK=%d\nPROC_ZOMBIE=%d\n", t, r, s, d, z}'
# 进程列表：etimes/times 都是纯数字，避免 start 字段含 "Jul 21" 多列错位
# comm 字段可能含空格，前 9 列用 \t 拼，comm 用换行做记录结束
ps -eo pid,user,pcpu,pmem,vsz,rss,stat,etimes,times,args --sort=-pcpu --no-headers 2>/dev/null \
  | head -200 \
  | awk '{
      out="PROC\t";
      for(i=1;i<=9;i++) out=out $i "\t";
      rest="";
      for(i=10;i<=NF;i++) rest=(i==10?$i:rest " " $i);
      print out rest
    }'
echo "===HOST==="
echo "HOSTNAME=$(hostname)"
echo "KERNEL=$(uname -r)"
echo "OS=$(. /etc/os-release 2>/dev/null && echo "$PRETTY_NAME" || uname -s)"
"""

    @staticmethod
    def _parse_kv(text: str, key: str, default: str = "0") -> str:
        """从 KEY=VALUE 行中取值"""
        m = re.search(rf"^{re.escape(key)}=(.+)$", text, re.MULTILINE)
        return m.group(1).strip() if m else default

    @classmethod
    def collect(cls, conn: SSHConnection) -> dict:
        """采集一次指标；返回 dict"""
        empty = {
            "cpu": 0.0, "cores": 1, "load1": 0, "load5": 0, "load15": 0,
            "uptime": 0, "hostname": "", "kernel": "", "os": "",
            "mem_total": 0, "mem_used": 0, "mem_percent": 0.0,
            "swap_total": 0, "swap_used": 0,
            "disks": [], "net": [],
            "proc_total": 0, "proc_running": 0, "proc_sleep": 0,
            "proc_disk": 0, "proc_zombie": 0, "processes": [],
            "ts": time.time(),
        }
        if not conn or not conn.connected:
            return empty
        code, out, err = conn.exec_command(cls._LINUX_METRICS_SCRIPT, timeout=10)
        if code != 0 or not out:
            empty["error"] = err or "采集失败"
            return empty

        result = dict(empty)
        result["hostname"] = cls._parse_kv(out, "HOSTNAME")
        result["kernel"] = cls._parse_kv(out, "KERNEL")
        result["os"] = cls._parse_kv(out, "OS")
        try:
            result["cpu"] = float(cls._parse_kv(out, "CPU_USAGE"))
        except ValueError:
            pass
        try:
            result["cores"] = int(cls._parse_kv(out, "CPU_CORES", "1"))
        except ValueError:
            pass
        load = cls._parse_kv(out, "LOAD", "0 0 0").split()
        try:
            result["load1"] = float(load[0])
            result["load5"] = float(load[1]) if len(load) > 1 else 0
            result["load15"] = float(load[2]) if len(load) > 2 else 0
        except (ValueError, IndexError):
            pass
        try:
            result["uptime"] = int(cls._parse_kv(out, "UPTIME"))
        except ValueError:
            pass

        try:
            mt = int(cls._parse_kv(out, "MEM_TOTAL"))
            mu = int(cls._parse_kv(out, "MEM_USED"))
            result["mem_total"] = mt
            result["mem_used"] = mu
            result["mem_percent"] = (mu / mt * 100) if mt > 0 else 0.0
            result["swap_total"] = int(cls._parse_kv(out, "SWAP_TOTAL"))
            result["swap_used"] = int(cls._parse_kv(out, "SWAP_USED"))
        except ValueError:
            pass

        result["disks"] = []
        for line in out.splitlines():
            if line.startswith("DISK|"):
                _, mount, total, used, percent = line.split("|", 4)
                try:
                    result["disks"].append({
                        "mount": mount, "total": int(total),
                        "used": int(used), "percent": int(percent.rstrip("%")),
                    })
                except ValueError:
                    continue

        result["net"] = []
        for line in out.splitlines():
            if line.startswith("NET|"):
                _, name, rx, tx = line.split("|", 3)
                try:
                    result["net"].append({
                        "iface": name, "rx": int(rx), "tx": int(tx),
                    })
                except ValueError:
                    continue

        # 进程统计
        try:
            result["proc_total"] = int(cls._parse_kv(out, "PROC_TOTAL"))
            result["proc_running"] = int(cls._parse_kv(out, "PROC_RUNNING"))
            result["proc_sleep"] = int(cls._parse_kv(out, "PROC_SLEEP"))
            result["proc_disk"] = int(cls._parse_kv(out, "PROC_DISK"))
            result["proc_zombie"] = int(cls._parse_kv(out, "PROC_ZOMBIE"))
        except ValueError:
            result["proc_total"] = 0
            result["proc_running"] = 0
            result["proc_sleep"] = 0
            result["proc_disk"] = 0
            result["proc_zombie"] = 0

        # 进程列表（远程用 \t 分隔前 9 列，第 10 列开始是 args/comm，可能含空格）
        result["processes"] = []
        for line in out.splitlines():
            if not line.startswith("PROC\t"):
                continue
            payload = line[5:]  # 去掉 "PROC\t" 前缀
            fields = payload.split("\t", 9)  # 只切前 9 次
            if len(fields) < 10:
                continue
            try:
                result["processes"].append({
                    "pid": int(fields[0]),
                    "user": fields[1],
                    "pcpu": float(fields[2]),
                    "pmem": float(fields[3]),
                    "vsz": int(fields[4]),
                    "rss": int(fields[5]),
                    "stat": fields[6],
                    # etimes: 自启动以来的秒数（整数）
                    "etime": int(fields[7]) if fields[7].isdigit() else 0,
                    # times: 累计 CPU 时间秒数
                    "time": int(fields[8]) if fields[8].isdigit() else 0,
                    "comm": fields[9].strip(),
                })
            except (ValueError, IndexError):
                continue

        return result

    @staticmethod
    def format_bytes(n: int) -> str:
        """人类可读字节数"""
        n = float(n)
        for unit in ("B", "KB", "MB", "GB", "TB", "PB"):
            if n < 1024:
                return f"{n:.1f}{unit}"
            n /= 1024
        return f"{n:.1f}EB"

    @staticmethod
    def format_uptime(seconds: int) -> str:
        seconds = int(seconds)
        d, rem = divmod(seconds, 86400)
        h, rem = divmod(rem, 3600)
        m, s = divmod(rem, 60)
        if d:
            return f"{d}天{h}小时"
        if h:
            return f"{h}小时{m}分"
        if m:
            return f"{m}分{s}秒"
        return f"{s}秒"