vps-manager/agent/collector.py

"""VPS Agent 采集模块

采集本机资源、服务器信息与基础安全状态。
依赖 psutil + 少量系统命令（ufw/iptables/sshd_config）。
"""

import glob
import platform
import socket
import subprocess
import time

import psutil


def collect_metrics() -> dict:
    """资源使用率（CPU/内存/磁盘/网络/负载/运行时长）"""
    cpu = psutil.cpu_percent(interval=1)
    mem = psutil.virtual_memory()
    disk = psutil.disk_usage("/")
    net = psutil.net_io_counters()
    try:
        load_1m = psutil.getloadavg()[0]
    except (AttributeError, OSError):
        load_1m = None
    return {
        "cpu_pct": round(cpu, 1),
        "mem_pct": round(mem.percent, 1),
        "disk_pct": round(disk.percent, 1),
        "net_in_mb": round(net.bytes_recv / 1024 / 1024, 2),
        "net_out_mb": round(net.bytes_sent / 1024 / 1024, 2),
        "load_1m": round(load_1m, 2) if load_1m is not None else None,
        "uptime_sec": int(time.time() - psutil.boot_time()),
    }


def _get_ip() -> str:
    """获取本机出口 IP（UDP 不实际发包；设超时避免无网环境阻塞）"""
    s = None
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(2)  # 无网络环境下 connect 可能阻塞，加超时保护
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]
    except Exception:
        return ""
    finally:
        if s is not None:
            try:
                s.close()
            except Exception:
                pass


def collect_server_info() -> dict:
    """服务器基础信息"""
    os_name = platform.system()
    try:
        with open("/etc/os-release") as f:
            for line in f:
                if line.startswith("PRETTY_NAME="):
                    os_name = line.split("=", 1)[1].strip().strip('"')
                    break
    except Exception:
        pass
    return {
        "hostname": socket.gethostname(),
        "os": os_name,
        "kernel": platform.release(),
        "cpu_cores": psutil.cpu_count(),
        "mem_total_gb": round(psutil.virtual_memory().total / 1024**3, 1),
        "disk_total_gb": round(psutil.disk_usage("/").total / 1024**3, 1),
        "public_ip": _get_ip(),
        "status": "online",
    }


SSHD_CONFIG = "/etc/ssh/sshd_config"
_SSHD_KEYS = ("permitrootlogin", "passwordauthentication")
# sshd 编译默认值：关键字未显式配置时实际生效的就是这些
_SSHD_DEFAULTS = {"permitrootlogin": "prohibit-password", "passwordauthentication": "yes"}


def _parse_sshd_config(path: str, values: dict, seen: set, depth: int = 0) -> bool:
    """按 sshd 语义解析配置：Include 就地展开、同一关键字「首次出现的值生效」

    返回是否成功打开该文件。Debian/Ubuntu 云镜像通常把加固项放在
    /etc/ssh/sshd_config.d/*.conf，而主配置顶部就有 Include；只读主文件会漏掉
    这些值，把已加固的机器误报成 warn，或反过来漏报。
    """
    real = path if path.startswith("/") else f"/etc/ssh/{path}"
    if depth > 3 or real in seen:
        return False
    seen.add(real)
    try:
        with open(real) as f:
            lines = f.readlines()
    except OSError:
        return False
    for line in lines:
        low = line.strip().lower()
        if not low or low.startswith("#"):
            continue
        parts = low.split()
        if parts[0] == "include" and len(parts) > 1:
            # Include 的相对路径由 sshd 按 /etc/ssh 解析，不能交给 glob 按 CWD 展开
            pattern = parts[1] if parts[1].startswith("/") else f"/etc/ssh/{parts[1]}"
            for inc in sorted(glob.glob(pattern)):
                _parse_sshd_config(inc, values, seen, depth + 1)
        elif parts[0] in _SSHD_KEYS and len(parts) > 1:
            values.setdefault(parts[0], parts[1])  # 首次出现生效
    return True


def _check_ssh_config() -> dict:
    """检查 SSH 配置（root 登录 / 密码登录），含 sshd_config.d 与默认值判定"""
    values: dict = {}
    if not _parse_sshd_config(SSHD_CONFIG, values, set()):
        return {"check_item": "ssh_config", "status": "unknown",
                "detail": "无法读取 sshd_config", "suggestion": None}

    permit_root = values.get("permitrootlogin", _SSHD_DEFAULTS["permitrootlogin"])
    password_auth = values.get("passwordauthentication", _SSHD_DEFAULTS["passwordauthentication"])
    # 未显式配置的关键字按 sshd 默认值判定，不能因为「没写」就当作已通过
    # （PasswordAuthentication 默认就是 yes，这类机器实际是开着密码登录的）
    implicit = [k for k in _SSHD_KEYS if k not in values]

    status = "pass"
    suggestions = []
    if permit_root == "yes":
        status = "warn"
        suggestions.append("设置 PermitRootLogin no 或 prohibit-password")
    if password_auth == "yes":
        status = "warn"
        suggestions.append("关闭密码登录 PasswordAuthentication no，改用密钥")
    detail = f"PermitRootLogin={permit_root}, PasswordAuthentication={password_auth}"
    if implicit:
        detail += f"（{'/'.join(implicit)} 未显式配置，取 sshd 默认值）"
    return {
        "check_item": "ssh_config",
        "status": status,
        "detail": detail,
        "suggestion": "；".join(suggestions) if suggestions else None,
    }


def _check_firewall() -> dict:
    """检查防火墙状态"""
    try:
        out = subprocess.run(["ufw", "status"], capture_output=True, text=True, timeout=5)
        if "Status: active" in out.stdout:
            return {"check_item": "firewall", "status": "pass",
                    "detail": "ufw 已启用", "suggestion": None}
        return {"check_item": "firewall", "status": "warn",
                "detail": "ufw 未启用", "suggestion": "启用防火墙 ufw enable"}
    except Exception:
        try:
            # 用 -S 而不是 -L：-L 即使零规则也会输出链头（Chain INPUT ... + 表头），
            # 拿 stdout 非空当「有规则」会永远判 pass
            out = subprocess.run(["iptables", "-S"], capture_output=True, text=True, timeout=5)
            lines = [ln.strip() for ln in out.stdout.splitlines() if ln.strip()]
            rules = [ln for ln in lines if ln.startswith("-A")]
            policy = next(
                (ln.split()[-1] for ln in lines if ln.startswith("-P INPUT")), "ACCEPT"
            )
            if rules or policy.upper() == "DROP":
                return {"check_item": "firewall", "status": "pass",
                        "detail": f"iptables {len(rules)} 条规则，INPUT 策略 {policy}",
                        "suggestion": None}
            return {"check_item": "firewall", "status": "warn",
                    "detail": f"iptables 无过滤规则（INPUT 策略 {policy}）",
                    "suggestion": "配置防火墙规则，或启用 ufw/nftables"}
        except Exception:
            pass
        return {"check_item": "firewall", "status": "unknown",
                "detail": "未检测到 ufw/iptables", "suggestion": None}


def _check_listening_ports() -> dict:
    """统计监听端口"""
    try:
        conns = psutil.net_connections(kind="inet")
        listening = sorted({c.laddr.port for c in conns if c.status == "LISTEN"})
        detail = f"开放端口 {len(listening)} 个: {','.join(map(str, listening[:15]))}"
        status = "pass" if len(listening) < 10 else "warn"
        suggestion = "端口偏多，建议关闭非必要端口" if status == "warn" else None
        return {"check_item": "listening_ports", "status": status,
                "detail": detail, "suggestion": suggestion}
    except Exception as e:  # noqa: BLE001
        return {"check_item": "listening_ports", "status": "unknown",
                "detail": str(e), "suggestion": None}


def collect_security() -> list:
    """基础安全检查项"""
    return [_check_ssh_config(), _check_firewall(), _check_listening_ports()]