"""SSH manager — similar to Amnezia-Web-Panel (Paramiko).""" from __future__ import annotations import io import logging import socket from dataclasses import dataclass import paramiko logger = logging.getLogger(__name__) @dataclass class SSHProbeResult: online: bool latency_ms: int | None = None error: str | None = None class SSHManager: """SSH connection + command execution on a remote VPS.""" def __init__( self, host: str, port: int = 22, username: str = "root", password: str | None = None, private_key: str | None = None, ): self.host = host self.port = int(port) self.username = username self.password = password or None self.private_key = private_key or None self.client: paramiko.SSHClient | None = None self._is_root = username == "root" def _load_pkey(self) -> paramiko.PKey: assert self.private_key is not None key_file = io.StringIO(self.private_key.strip()) errors: list[str] = [] for key_cls in (paramiko.Ed25519Key, paramiko.RSAKey, paramiko.ECDSAKey): key_file.seek(0) try: return key_cls.from_private_key(key_file) except Exception as exc: # noqa: BLE001 errors.append(f"{key_cls.__name__}: {exc}") raise ValueError("Не удалось прочитать приватный ключ (" + "; ".join(errors[:2]) + ")") def connect(self, timeout: int = 15) -> bool: if not self.password and not self.private_key: raise ValueError("Нужен пароль или приватный ключ SSH") self.client = paramiko.SSHClient() self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) kwargs: dict = { "hostname": self.host, "port": self.port, "username": self.username, "timeout": timeout, "banner_timeout": timeout, "auth_timeout": timeout, "allow_agent": False, "look_for_keys": False, } if self.private_key: kwargs["pkey"] = self._load_pkey() else: kwargs["password"] = self.password self.client.connect(**kwargs) return True def disconnect(self) -> None: if self.client: self.client.close() self.client = None def run_command(self, command: str, timeout: int = 60) -> tuple[str, str, int]: if not self.client: raise ConnectionError("Not connected to server") logger.info("SSH %s: %s", self.host, command[:120]) _stdin, stdout, stderr = self.client.exec_command(command, timeout=timeout) stdout.channel.settimeout(timeout) stderr.channel.settimeout(timeout) try: exit_code = stdout.channel.recv_exit_status() out = stdout.read().decode("utf-8", errors="replace").strip() err = stderr.read().decode("utf-8", errors="replace").strip() except Exception as exc: # noqa: BLE001 return "", str(exc), -1 return out, err, exit_code def run_sudo_command(self, command: str, timeout: int = 60) -> tuple[str, str, int]: clean = command.strip() if clean.startswith("sudo "): clean = clean[5:] if self._is_root: return self.run_command(clean, timeout=timeout) if self.password: escaped = self.password.replace("'", "'\\''") full = f"echo '{escaped}' | sudo -S -p '' {clean}" else: full = f"sudo {clean}" return self.run_command(full, timeout=timeout) def test_connection(self) -> str: out, err, code = self.run_command( "uname -sr && cat /etc/os-release 2>/dev/null | head -2" ) if code != 0 and not out: raise ConnectionError(err or f"SSH test failed (code {code})") return out def __enter__(self) -> SSHManager: self.connect() return self def __exit__(self, *args) -> None: self.disconnect() def probe_tcp(host: str, port: int, timeout: float = 2.0) -> SSHProbeResult: """Non-blocking-ish TCP connect probe (run via asyncio.to_thread).""" import time started = time.perf_counter() try: with socket.create_connection((host, int(port)), timeout=timeout): latency = int((time.perf_counter() - started) * 1000) return SSHProbeResult(online=True, latency_ms=latency) except OSError as exc: return SSHProbeResult(online=False, error=str(exc)) def verify_ssh( host: str, port: int, username: str, password: str | None = None, private_key: str | None = None, ) -> str: """Connect, run test command, disconnect. Returns server_info text.""" ssh = SSHManager(host, port, username, password, private_key) try: ssh.connect() return ssh.test_connection() finally: ssh.disconnect()