Nexus/server/infrastructure/ssh/asyncssh_pool.py

"""Nexus — asyncssh Connection Pool (Reference Counting Mode)

ADR-002: BFF mode with asyncssh embedded (single process, full Python).
ADR-004: Reference counting connection pool — reuse connections, reduce handshakes.

W1: Extracted from ssh_direct.py, provides asyncssh-based connection pooling.
W2: Replaces paramiko pool.py with asyncssh throughout the codebase.
"""

import asyncio
import logging
import secrets
import shlex
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import AsyncIterator, Dict, Optional

import asyncssh
from server.domain.models import Server
from server.utils.files_elevation import FilesElevation
from server.utils.posix_paths import posix_dirname
from server.infrastructure.database.crypto import decrypt_value
from server.utils.sync_error_message import translate_ssh_error_message

logger = logging.getLogger("nexus.asyncssh_pool")


def _localized_ssh_error(exc: Exception, *, max_len: int = 200) -> str:
    raw = str(exc)[:max_len]
    return translate_ssh_error_message(raw) or raw


async def try_ssh_login(
    host: str,
    port: int,
    username: str,
    *,
    password: str | None = None,
    private_key: str | None = None,
    timeout: float = 8.0,
) -> tuple[bool, str]:
    """One-shot SSH login probe (not pooled). Returns (ok, error_message)."""
    if not password and not private_key:
        return False, "未提供密码或私钥"

    connect_kwargs: dict = {
        "host": host,
        "port": port or 22,
        "username": username or "root",
        "connect_timeout": timeout,
        # 凭据轮询探测未知主机：跳过主机密钥校验（运维平台标准；持久连接见 SSH_STRICT_HOST_CHECKING）
        "known_hosts": None,
    }

    if password:
        connect_kwargs["password"] = password
    elif private_key:
        connect_kwargs["client_keys"] = [asyncssh.import_private_key(private_key)]

    conn: asyncssh.SSHClientConnection | None = None
    try:
        conn = await asyncssh.connect(**connect_kwargs)
        return True, ""
    except asyncssh.Error as e:
        return False, _localized_ssh_error(e)
    except Exception as e:
        return False, _localized_ssh_error(e)
    finally:
        if conn is not None:
            try:
                conn.close()
            except Exception:  # noqa: S110 — best-effort cleanup
                pass


@dataclass
class PooledConnection:
    """A tracked asyncssh connection with reference counting"""
    conn: asyncssh.SSHClientConnection
    ref_count: int = 0
    created_at: float = field(default_factory=time.monotonic)
    last_used: float = field(default_factory=time.monotonic)
    server_id: int = 0

    @property
    def is_idle(self) -> bool:
        return self.ref_count == 0

    @property
    def idle_seconds(self) -> float:
        return time.monotonic() - self.last_used if self.is_idle else 0


class AsyncSSHPool:
    """asyncssh connection pool with reference counting (ADR-004)

    Features:
    - Reference counting: multiple operations can share one SSH connection
    - Auto-cleanup: idle connections closed after IDLE_TIMEOUT
    - Per-server pooling: one connection per server at a time
    - Thread-safe: uses asyncio.Lock for concurrent access
    """

    # Pool configuration
    MAX_CONNECTIONS = 100       # Global max connections
    IDLE_TIMEOUT = 300          # 5 minutes idle → close
    CONNECT_TIMEOUT = 10        # 10 seconds to establish connection
    CLEANUP_INTERVAL = 60       # Check for idle connections every 60s

    def __init__(self):
        self._pool: Dict[int, PooledConnection] = {}  # server_id → connection
        self._lock = asyncio.Lock()
        self._cleanup_task: Optional[asyncio.Task] = None

    async def start(self):
        """Start the connection pool and cleanup task"""
        self._cleanup_task = asyncio.create_task(self._cleanup_loop(), name="asyncssh_cleanup")
        logger.info(f"AsyncSSH pool started (max={self.MAX_CONNECTIONS}, idle_timeout={self.IDLE_TIMEOUT}s)")

    async def stop(self):
        """Stop the pool and close all connections"""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        async with self._lock:
            for _server_id, pooled in self._pool.items():
                try:
                    pooled.conn.close()
                except Exception:  # noqa: S110 — best-effort cleanup
                    pass
            self._pool.clear()
        logger.info("AsyncSSH pool stopped — all connections closed")

    async def acquire(self, server: Server) -> asyncssh.SSHClientConnection:
        """Acquire a connection to a server (reference counting).

        If an idle connection exists for this server, reuse it.
        Otherwise, create a new connection.
        SSH handshake is done outside the lock to avoid blocking other pool operations.
        """
        from server.infrastructure.ssh.resolve_server_auth import materialize_server_ssh_auth

        await materialize_server_ssh_auth(server)
        async with self._lock:
            # Check for existing idle connection
            pooled = self._pool.get(server.id)
            if pooled and not pooled.conn.is_closed():
                pooled.ref_count += 1
                pooled.last_used = time.monotonic()
                logger.debug(f"Reused SSH connection: server={server.id}, refs={pooled.ref_count}")
                return pooled.conn

            # Need to create — evict if at capacity
            if len(self._pool) >= self.MAX_CONNECTIONS:
                evicted = await self._evict_one()
                if not evicted:
                    # All connections are busy — cannot exceed MAX_CONNECTIONS
                    raise ConnectionError(
                        f"SSH pool at capacity ({self.MAX_CONNECTIONS}) and all connections busy. "
                        f"Try again later."
                    )

        # Create connection outside lock (handshake can take seconds)
        conn = await self._create_connection(server)

        async with self._lock:
            # Re-check: another coroutine may have created one while we were unlocked
            pooled = self._pool.get(server.id)
            if pooled and not pooled.conn.is_closed():
                # Close our new connection, reuse the existing one
                try:
                    conn.close()
                except Exception:  # noqa: S110 — best-effort cleanup
                    pass
                pooled.ref_count += 1
                pooled.last_used = time.monotonic()
                return pooled.conn

            pooled = PooledConnection(
                conn=conn,
                ref_count=1,
                server_id=server.id,
            )
            self._pool[server.id] = pooled
            logger.debug(f"New SSH connection: server={server.id}, pool_size={len(self._pool)}")
            return conn

    async def release(self, server_id: int):
        """Release a connection (decrement reference count).

        If ref_count reaches 0, the connection stays in pool for reuse.
        It will be cleaned up by the idle timeout task.
        """
        async with self._lock:
            pooled = self._pool.get(server_id)
            if pooled:
                pooled.ref_count = max(0, pooled.ref_count - 1)
                pooled.last_used = time.monotonic()
                logger.debug(f"Released SSH connection: server={server_id}, refs={pooled.ref_count}")

    async def close_connection(self, server_id: int):
        """Force-close a connection to a specific server"""
        async with self._lock:
            pooled = self._pool.pop(server_id, None)
            if pooled:
                try:
                    pooled.conn.close()
                except Exception:  # noqa: S110 — best-effort cleanup
                    pass
                logger.debug(f"Force-closed SSH connection: server={server_id}")

    async def _create_connection(self, server: Server) -> asyncssh.SSHClientConnection:
        """Create a new asyncssh connection to a server"""
        from server.config import settings

        connect_kwargs = {
            "host": server.domain,
            "port": server.port or 22,
            "username": server.username or "root",
            "connect_timeout": self.CONNECT_TIMEOUT,
        }

        # Host key verification: respect SSH_STRICT_HOST_CHECKING setting
        # Default is "true" — reject unknown hosts for security
        if settings.SSH_STRICT_HOST_CHECKING.lower() in ("false", "0", "no"):
            connect_kwargs["known_hosts"] = None  # Skip verification (insecure)
        # If strict mode, asyncssh will use default known_hosts behavior
        # (rejects unknown hosts — caller should configure ~/.ssh/known_hosts)

        if server.auth_method == "password" and server.password:
            password = decrypt_value(server.password)
            connect_kwargs["password"] = password
        elif server.ssh_key_configured and server.ssh_key_private:
            key_content = decrypt_value(server.ssh_key_private)
            connect_kwargs["client_keys"] = [asyncssh.import_private_key(key_content)]
        elif server.ssh_key_path:
            connect_kwargs["client_keys"] = [server.ssh_key_path]
        else:
            raise ValueError(f"服务器 {server.id}（{server.name}）没有可用的 SSH 凭据")

        try:
            conn = await asyncssh.connect(**connect_kwargs)
            return conn
        except asyncssh.Error as e:
            msg = _localized_ssh_error(e)
            raise ConnectionError(
                f"SSH 连接失败 {server.domain}:{server.port or 22}：{msg}"
            ) from e

    async def _evict_one(self) -> bool:
        """Evict the oldest idle connection from the pool.

        Returns True if an idle connection was evicted, False if all are busy.
        """
        oldest_idle = None
        oldest_time = float("inf")

        for server_id, pooled in self._pool.items():
            if pooled.is_idle and pooled.last_used < oldest_time:
                oldest_idle = server_id
                oldest_time = pooled.last_used

        if oldest_idle is not None:
            pooled = self._pool.pop(oldest_idle)
            try:
                pooled.conn.close()
            except Exception:  # noqa: S110 — best-effort cleanup
                pass
            logger.debug(f"Evicted idle SSH connection: server={oldest_idle}")
            return True
        return False

    async def _cleanup_loop(self):
        """Periodically clean up idle connections"""
        while True:
            await asyncio.sleep(self.CLEANUP_INTERVAL)
            try:
                async with self._lock:
                    to_remove = []
                    for server_id, pooled in self._pool.items():
                        if pooled.is_idle and pooled.idle_seconds > self.IDLE_TIMEOUT:
                            to_remove.append(server_id)
                        elif not pooled.is_idle and pooled.conn.is_closed():
                            # Connection died while in use — remove from pool
                            to_remove.append(server_id)

                    for server_id in to_remove:
                        pooled = self._pool.pop(server_id)
                        try:
                            pooled.conn.close()
                        except Exception:  # noqa: S110 — best-effort cleanup
                            pass

                    if to_remove:
                        logger.info(f"SSH pool cleanup: removed {len(to_remove)} idle/dead connections")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"SSH pool cleanup error: {e}")

    @property
    def stats(self) -> dict:
        """Get pool statistics"""
        total = len(self._pool)
        active = sum(1 for p in self._pool.values() if not p.is_idle)
        idle = total - active
        return {"total": total, "active": active, "idle": idle}


# Global pool instance
ssh_pool = AsyncSSHPool()


@asynccontextmanager
async def sftp_session(conn: asyncssh.SSHClientConnection) -> AsyncIterator[asyncssh.SFTPClient]:
    """Open an SFTP client on a pooled SSH connection (asyncssh 2.x: start_sftp_client)."""
    async with await conn.start_sftp_client() as sftp:
        yield sftp


async def sftp_write_bytes(sftp: asyncssh.SFTPClient, remote_path: str, data: bytes) -> None:
    """Write in-memory bytes to a remote file (asyncssh put() only accepts local paths)."""
    async with sftp.open(remote_path, "wb") as remote_file:
        await remote_file.write(data)


def _sftp_permission_denied(exc: Exception) -> bool:
    if isinstance(exc, asyncssh.PermissionDenied):
        return True
    if isinstance(exc, asyncssh.SFTPError):
        # SSH_FX_PERMISSION_DENIED = 3
        if getattr(exc, "code", None) == 3:
            return True
    return "permission denied" in str(exc).lower()


async def _run_tee_stdin(
    conn: asyncssh.SSHClientConnection,
    command: str,
    data: bytes,
) -> asyncssh.SSHCompletedProcess:
    """Feed raw bytes to stdin; default SSH session encoding expects str, not bytes."""
    return await conn.run(command, input=data, encoding=None, check=False)


async def _ensure_remote_parent_dir(
    conn: asyncssh.SSHClientConnection,
    remote_path: str,
    *,
    use_sudo: bool,
) -> None:
    parent = posix_dirname(remote_path)
    if not parent or parent == "/":
        return
    prefix = "sudo -n " if use_sudo else ""
    await conn.run(f"{prefix}mkdir -p {shlex.quote(parent)}", check=False)


async def _ssh_write_bytes_via_tee(
    conn: asyncssh.SSHClientConnection,
    remote_path: str,
    data: bytes,
    *,
    use_sudo: bool,
) -> None:
    """Shell write fallback when SFTP cannot open the file (e.g. file not writable but directory is)."""
    prefix = "sudo -n " if use_sudo else ""
    await _ensure_remote_parent_dir(conn, remote_path, use_sudo=use_sudo)

    direct = await _run_tee_stdin(
        conn,
        f"{prefix}tee {shlex.quote(remote_path)}",
        data,
    )
    if direct.exit_status == 0:
        return

    tmp_path = f"/tmp/nexus-write-{secrets.token_hex(12)}"
    staged = await _run_tee_stdin(
        conn,
        f"{prefix}tee {shlex.quote(tmp_path)}",
        data,
    )
    if staged.exit_status != 0:
        stderr = (staged.stderr or staged.stdout or direct.stderr or "")[:500]
        raise PermissionError(stderr or "tee 写入失败")

    moved = await conn.run(
        f"{prefix}mv -f {shlex.quote(tmp_path)} {shlex.quote(remote_path)}",
        check=False,
    )
    if moved.exit_status != 0:
        await conn.run(f"{prefix}rm -f {shlex.quote(tmp_path)}", check=False)
        stderr = (moved.stderr or "")[:500]
        raise PermissionError(stderr or "无法覆盖目标文件（目录可能不可写）")


def _sudo_attempt_order(elevation: FilesElevation | None) -> tuple[bool, ...]:
    mode = elevation if elevation is not None else FilesElevation.AUTO
    if mode == FilesElevation.OFF:
        return (False,)
    if mode == FilesElevation.ALWAYS:
        return (True,)
    return (False, True)


async def remote_write_bytes(
    conn: asyncssh.SSHClientConnection,
    remote_path: str,
    data: bytes,
    *,
    elevation: FilesElevation | None = None,
) -> None:
    """Write bytes to a remote path: SFTP first, then SSH tee / tmp+mv (incl. optional sudo -n)."""
    eff = elevation if elevation is not None else FilesElevation.AUTO
    skip_sftp = eff == FilesElevation.ALWAYS

    if not skip_sftp:
        try:
            parent = posix_dirname(remote_path)
            if parent and parent != "/":
                await _ensure_remote_parent_dir(conn, remote_path, use_sudo=False)
            async with await conn.start_sftp_client() as sftp:
                await sftp_write_bytes(sftp, remote_path, data)
            return
        except (asyncssh.PermissionDenied, asyncssh.SFTPError) as exc:
            if not _sftp_permission_denied(exc):
                raise

    last_error: Exception | None = None
    for use_sudo in _sudo_attempt_order(eff):
        try:
            await _ssh_write_bytes_via_tee(conn, remote_path, data, use_sudo=use_sudo)
            return
        except PermissionError as exc:
            last_error = exc

    reason = str(last_error) if last_error else "Permission denied"
    raise asyncssh.SFTPError(3, reason)


async def exec_ssh_command(
    server: Server,
    command: str,
    timeout: int = 30,
    *,
    login_shell: bool = True,
) -> dict:
    """Execute a command via SSH using the asyncssh connection pool.

    W2: Replaces paramiko-based exec_command from pool.py.
    Returns structured result dict with stdout/stderr/exit_code.

    By default commands run in a login shell (``bash -lc``) so remote
    ``/etc/profile`` and ``~/.profile`` apply — matching interactive SSH.
    Pass ``login_shell=False`` only for rare cases that must avoid profile side
    effects (internal pool helpers should call ``conn.run`` directly).
    """
    from server.infrastructure.ssh.shell_wrap import wrap_login_shell

    remote_cmd = wrap_login_shell(command) if login_shell else command
    conn = None
    try:
        conn = await ssh_pool.acquire(server)
        result = await asyncio.wait_for(
            conn.run(remote_cmd, timeout=timeout),
            timeout=timeout + 5,
        )
        await ssh_pool.release(server.id)
        return {
            "status": "success" if result.exit_status == 0 else "failed",
            "stdout": result.stdout[:10000],
            "stderr": result.stderr[:10000],
            "exit_code": result.exit_status,
        }
    except asyncio.TimeoutError:
        # Timed-out connection may be in an unknown state — evict so it is NOT
        # returned to the pool and potentially reused for a future command.
        if conn is not None:
            await ssh_pool.close_connection(server.id)
        return {
            "status": "timeout",
            "stdout": "",
            "stderr": f"命令执行超时（{timeout} 秒）",
            "exit_code": -1,
        }
    except ConnectionError as e:
        if conn is not None:
            await ssh_pool.close_connection(server.id)
        return {
            "status": "connection_error",
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1,
        }
    except Exception as e:
        if conn is not None:
            await ssh_pool.release(server.id)
        return {
            "status": "error",
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1,
        }