debian-manager/app/core/ssh.py

"""SSH-слой: подключения, выполнение команд, sudo, потоковый вывод."""
from __future__ import annotations

import asyncio
import shlex
from dataclasses import dataclass
from typing import AsyncIterator, Optional

import asyncssh

from app.db.models import Server

CONNECT_TIMEOUT = 15


@dataclass
class CmdResult:
    stdout: str
    stderr: str
    returncode: int

    @property
    def ok(self) -> bool:
        return self.returncode == 0


class SSHClient:
    """Одно SSH-подключение к серверу с помощниками sudo."""

    def __init__(self, server: Server):
        self.server = server
        self._conn: Optional[asyncssh.SSHClientConnection] = None

    async def __aenter__(self) -> "SSHClient":
        await self.connect()
        return self

    async def __aexit__(self, *exc) -> None:
        self.close()

    async def connect(self) -> None:
        kwargs: dict = dict(
            host=self.server.host,
            port=self.server.port,
            username=self.server.username,
            known_hosts=None,
            connect_timeout=CONNECT_TIMEOUT,
        )
        if self.server.auth_type == "key" and self.server.key_path:
            kwargs["client_keys"] = [self.server.key_path]
            if self.server.password:
                kwargs["passphrase"] = self.server.password
        else:
            kwargs["password"] = self.server.password
        self._conn = await asyncssh.connect(**kwargs)

    def close(self) -> None:
        if self._conn:
            self._conn.close()
            self._conn = None

    @property
    def conn(self) -> asyncssh.SSHClientConnection:
        if not self._conn:
            raise RuntimeError("SSH not connected")
        return self._conn

    # ---------- выполнение ----------

    async def run(self, cmd: str, timeout: int = 120) -> CmdResult:
        r = await self.conn.run(cmd, timeout=timeout)
        return CmdResult(r.stdout or "", r.stderr or "", r.returncode)

    def sudo_wrap(self, cmd: str) -> str:
        """Обернуть команду в sudo -S с передачей пароля по stdin."""
        pw = self.server.sudo_password or self.server.password or ""
        if self.server.username == "root":
            return cmd
        return f"echo {shlex.quote(pw)} | sudo -S -p '' bash -c {shlex.quote(cmd)}"

    async def sudo(self, cmd: str, timeout: int = 600) -> CmdResult:
        return await self.run(self.sudo_wrap(cmd), timeout=timeout)

    async def run_stream(self, cmd: str, use_sudo: bool = False) -> AsyncIterator[str]:
        """Строки вывода по мере поступления (для лог-панели)."""
        full = self.sudo_wrap(cmd) if use_sudo else cmd
        async with self.conn.create_process(full, term_type="dumb") as proc:
            async for line in proc.stdout:
                yield line.rstrip("\n")
            await proc.wait()

    async def is_alive(self) -> bool:
        try:
            r = await self.run("echo ok", timeout=10)
            return r.ok
        except Exception:
            return False


async def check_server(server: Server) -> dict:
    """Быстрый health-check: онлайн, ОС, uptime."""
    try:
        async with SSHClient(server) as c:
            r = await c.run(
                "hostname; . /etc/os-release && echo $PRETTY_NAME; "
                "uptime -p; nproc; free -m | awk 'NR==2{print $2}'",
                timeout=20,
            )
            lines = (r.stdout or "").splitlines()
            return {
                "online": True,
                "hostname": lines[0] if len(lines) > 0 else "",
                "os": lines[1] if len(lines) > 1 else "",
                "uptime": lines[2] if len(lines) > 2 else "",
                "cpu_cores": lines[3] if len(lines) > 3 else "?",
                "mem_mb": lines[4] if len(lines) > 4 else "?",
            }
    except Exception as e:
        return {"online": False, "error": str(e)}


async def wait_online(server: Server, timeout_s: int = 300, interval: int = 5) -> bool:
    """Ждать возврата сервера после reboot."""
    deadline = asyncio.get_event_loop().time() + timeout_s
    while asyncio.get_event_loop().time() < deadline:
        info = await check_server(server)
        if info.get("online"):
            return True
        await asyncio.sleep(interval)
    return False