vpn/app/services/ssh.py

"""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 upload_file(self, content: str, remote_path: str) -> None:
        if not self.client:
            raise ConnectionError("Not connected to server")
        content = content.replace("\r\n", "\n")
        sftp = self.client.open_sftp()
        try:
            with sftp.file(remote_path, "w") as f:
                f.write(content)
        finally:
            sftp.close()

    def upload_file_sudo(self, content: str, remote_path: str) -> None:
        import hashlib

        tmp = f"/tmp/_vpnpanel_{hashlib.md5(remote_path.encode()).hexdigest()[:8]}"
        self.upload_file(content, tmp)
        self.run_sudo_command(f"mv {tmp} {remote_path}")
        self.run_sudo_command(f"chmod 644 {remote_path}")

    def run_sudo_script(self, script: str, timeout: int = 180) -> tuple[str, str, int]:
        import hashlib

        script = script.replace("\r\n", "\n")
        tmp = f"/tmp/_vpnpanel_script_{hashlib.md5(script.encode()).hexdigest()[:8]}.sh"
        self.upload_file(script, tmp)
        if self._is_root:
            out = self.run_command(f"bash {tmp}; rm -f {tmp}", timeout=timeout)
        elif self.password:
            escaped = self.password.replace("'", "'\\''")
            out = self.run_command(
                f"echo '{escaped}' | sudo -S -p '' bash {tmp}; rm -f {tmp}",
                timeout=timeout,
            )
        else:
            out = self.run_command(f"sudo bash {tmp}; rm -f {tmp}", timeout=timeout)
        return out

    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()