debian-manager/app/core/agent.py

"""Агент-режим: установка агента на сервер, WS-сервер, приём метрик."""
from __future__ import annotations

import asyncio
import json
import secrets
import time
from typing import Callable, Optional

from app.core.ssh import SSHClient
from app.db.database import get_session
from app.db.models import MetricSnapshot, Server, Setting

# последние метрики по server_id: {"cpu": .., "mem": .., "disk": .., "ts": ..}
agent_metrics: dict[int, dict] = {}
_http_server = None

AGENT_VERSION = "0.1"

AGENT_SCRIPT = r'''#!/usr/bin/env python3
"""Debian Manager agent: метрики по WebSocket. Только stdlib."""
import json, socket, ssl, sys, time, urllib.request, os
try:
    import psutil
except ImportError:
    psutil = None

SERVER_WS = os.environ.get("DM_SERVER", "ws://127.0.0.1:9101")
TOKEN = os.environ.get("DM_TOKEN", "")
SERVER_ID = int(os.environ.get("DM_SERVER_ID", "0"))
INTERVAL = float(os.environ.get("DM_INTERVAL", "5"))

def metrics():
    cpu = 0.0
    if psutil:
        cpu = psutil.cpu_percent(interval=None)
        vm = psutil.virtual_memory()
        mem_used, mem_total = vm.used // 1048576, vm.total // 1048576
        du = psutil.disk_usage("/")
        disk = du.percent
    else:
        with open("/proc/stat") as f:
            parts = f.readline().split()
            idle = int(parts[4]) + int(parts[5])
            total = sum(int(x) for x in parts[1:8])
            cpu = round((1 - idle / total) * 100, 1) if total else 0
        with open("/proc/meminfo") as f:
            mi = {}
            for line in f:
                k, v = line.split(":")
                mi[k] = int(v.split()[0])
            mem_total = mi["MemTotal"] // 1024
            mem_used = (mi["MemTotal"] - mi.get("MemAvailable", mi["MemFree"])) // 1024
        st = os.statvfs("/")
        disk = round((st.f_blocks - st.f_bavail) * 100 / st.f_blocks, 1)
    with open("/proc/loadavg") as f:
        load = float(f.read().split()[0])
    return {"cpu": cpu, "mem_used": mem_used, "mem_total": mem_total,
            "mem_pct": round(mem_used * 100 / mem_total, 1) if mem_total else 0,
            "disk": disk, "load": load, "ts": time.time()}

def main():
    # минимальный WebSocket-клиент на stdlib (только текстовые фреймы, без маскировки сложная;
    # для простоты используем HTTP POST каждые INTERVAL сек — надёжнее и без библиотек)
    url = SERVER_WS.replace("ws://", "http://").replace("wss://", "https://") + "/metrics"
    while True:
        try:
            data = json.dumps({"server_id": SERVER_ID, "token": TOKEN,
                               "metrics": metrics()}).encode()
            req = urllib.request.Request(url, data=data,
                                         headers={"Content-Type": "application/json"})
            urllib.request.urlopen(req, timeout=10)
        except Exception as e:
            print("send error:", e, file=sys.stderr)
        time.sleep(INTERVAL)

if __name__ == "__main__":
    main()
'''


def _get(key: str) -> str:
    with get_session() as s:
        row = s.get(Setting, key)
        return row.value if row else ""


def _set(key: str, value: str) -> None:
    with get_session() as s:
        row = s.get(Setting, key) or Setting(key=key)
        row.value = value
        s.add(row)
        s.commit()


def get_agent_port() -> int:
    return int(_get("agent_port") or "9102")


def get_or_create_token(server_id: int) -> str:
    key = f"agent_token_{server_id}"
    token = _get(key)
    if not token:
        token = secrets.token_hex(16)
        _set(key, token)
    return token


async def install_agent(srv: Server, app_host: str, log=None) -> bool:
    """Установить агента на сервер через SSH."""
    def _log(m):
        if log:
            log(m)
    token = get_or_create_token(srv.id)
    port = get_agent_port()
    env = (f"DM_SERVER=http://{app_host}:{port} DM_TOKEN={token} "
           f"DM_SERVER_ID={srv.id} DM_INTERVAL=5")
    script_b64 = __import__("base64").b64encode(AGENT_SCRIPT.encode()).decode()

    unit = f"""[Unit]
Description=Debian Manager Agent
After=network.target

[Service]
Type=simple
Environment="{env.replace('"', chr(92)+'"')}"
ExecStart=/usr/bin/python3 /opt/dm-agent/agent.py
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target
"""
    unit_b64 = __import__("base64").b64encode(unit.encode()).decode()

    async with SSHClient(srv) as c:
        _log("создаю /opt/dm-agent")
        await c.sudo("mkdir -p /opt/dm-agent")
        _log("записываю agent.py")
        await c.sudo(f"echo {script_b64} | base64 -d > /opt/dm-agent/agent.py")
        await c.sudo(f"echo {unit_b64} | base64 -d > /etc/systemd/system/dm-agent.service")
        _log("запускаю systemd unit")
        r = await c.sudo("systemctl daemon-reload && systemctl enable --now dm-agent")
        if not r.ok:
            _log(f"ошибка: {r.stderr}")
            return False
        _log("агент установлен и запущен")
        return True


async def uninstall_agent(srv: Server, log=None) -> None:
    def _log(m):
        if log:
            log(m)
    async with SSHClient(srv) as c:
        await c.sudo("systemctl disable --now dm-agent 2>/dev/null; "
                     "rm -f /etc/systemd/system/dm-agent.service; "
                     "rm -rf /opt/dm-agent; systemctl daemon-reload")
    _set(f"agent_token_{srv.id}", "")
    _log("агент удалён")


async def agent_status(srv: Server) -> str:
    async with SSHClient(srv) as c:
        r = await c.sudo("systemctl is-active dm-agent 2>/dev/null")
        return r.stdout.strip() or "не установлен"


# ---------- HTTP-приёмник метрик ----------

async def _metrics_handler(reader, writer):
    try:
        req = await reader.readuntil(b"\r\n\r\n")
        head = req.decode(errors="replace")
        length = 0
        for line in head.splitlines():
            if line.lower().startswith("content-length:"):
                length = int(line.split(":", 1)[1].strip())
        body = await reader.readexactly(length) if length else b"{}"
        data = json.loads(body)
        sid = int(data.get("server_id", 0))
        if data.get("token") != _get(f"agent_token_{sid}"):
            raise ValueError("bad token")
        m = data.get("metrics", {})
        agent_metrics[sid] = m
        with get_session() as s:
            s.add(MetricSnapshot(
                server_id=sid,
                cpu_percent=m.get("cpu", 0),
                mem_percent=m.get("mem_pct", 0),
                mem_used_mb=m.get("mem_used", 0),
                mem_total_mb=m.get("mem_total", 0),
                disk_percent=m.get("disk", 0),
                load1=m.get("load", 0),
            ))
            s.commit()
        resp = b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok"
    except Exception:
        resp = b"HTTP/1.1 400 Bad\r\nContent-Length: 0\r\n\r\n"
    writer.write(resp)
    await writer.drain()
    try:
        writer.close()
    except Exception:
        pass


async def start(port: Optional[int] = None) -> None:
    global _http_server
    if _http_server:
        return
    port = port or get_agent_port()
    _http_server = await asyncio.start_server(_metrics_handler, "0.0.0.0", port)


def is_running() -> bool:
    return _http_server is not None