"""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