95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
|
"""Nexus — SSH Connection Pool (Unified, replacing 4+ duplicated implementations)
|
||
|
|
|
||
|
|
Provides:
|
||
|
|
- SSHConfig dataclass for connection parameters
|
||
|
|
- ssh_connection() async context manager
|
||
|
|
- create_ssh_client() factory with SSH Key support
|
||
|
|
- Batch operations via async pool
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import contextlib
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
import paramiko
|
||
|
|
from server.domain.models import Server
|
||
|
|
from server.infrastructure.database.crypto import decrypt_value
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class SSHConfig:
|
||
|
|
"""SSH connection configuration"""
|
||
|
|
host: str
|
||
|
|
port: int = 22
|
||
|
|
username: str = "root"
|
||
|
|
auth_method: str = "key" # key/password
|
||
|
|
password: Optional[str] = None
|
||
|
|
key_path: Optional[str] = None
|
||
|
|
key_content: Optional[str] = None # For in-memory key (from DB encrypted field)
|
||
|
|
strict_host_checking: str = "true"
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_server(cls, server: Server) -> "SSHConfig":
|
||
|
|
"""Build SSHConfig from Server model (with decrypted password)"""
|
||
|
|
return cls(
|
||
|
|
host=server.domain,
|
||
|
|
port=server.port,
|
||
|
|
username=server.username,
|
||
|
|
auth_method=server.auth_method,
|
||
|
|
password=server.decrypted_password if server.auth_method == "password" else None,
|
||
|
|
key_path=server.ssh_key_path,
|
||
|
|
key_content=decrypt_value(server.ssh_key_private or "") if server.ssh_key_configured else None,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def create_ssh_client(config: SSHConfig) -> paramiko.SSHClient:
|
||
|
|
"""Create SSH client with key or password authentication"""
|
||
|
|
client = paramiko.SSHClient()
|
||
|
|
|
||
|
|
if config.strict_host_checking == "true":
|
||
|
|
client.set_missing_host_key_policy(paramiko.WarningPolicy())
|
||
|
|
else:
|
||
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
|
|
|
||
|
|
if config.auth_method == "key" and config.key_content:
|
||
|
|
# Use in-memory key from encrypted DB field
|
||
|
|
import io
|
||
|
|
key = paramiko.Ed25519Key.from_private_key(io.StringIO(config.key_content))
|
||
|
|
client.connect(config.host, port=config.port, username=config.username, pkey=key)
|
||
|
|
elif config.auth_method == "key" and config.key_path:
|
||
|
|
client.connect(config.host, port=config.port, username=config.username, key_filename=config.key_path)
|
||
|
|
elif config.auth_method == "password" and config.password:
|
||
|
|
client.connect(config.host, port=config.port, username=config.username, password=config.password)
|
||
|
|
else:
|
||
|
|
raise ValueError(f"No valid SSH credentials for {config.host}")
|
||
|
|
|
||
|
|
return client
|
||
|
|
|
||
|
|
|
||
|
|
@contextlib.asynccontextmanager
|
||
|
|
async def ssh_connection(config: SSHConfig):
|
||
|
|
"""Async context manager for SSH connections (creates and cleans up)"""
|
||
|
|
client = await asyncio.to_thread(create_ssh_client, config)
|
||
|
|
try:
|
||
|
|
yield client
|
||
|
|
finally:
|
||
|
|
client.close()
|
||
|
|
|
||
|
|
|
||
|
|
async def exec_command(config: SSHConfig, command: str, timeout: int = 30) -> dict:
|
||
|
|
"""Execute a command via SSH and return structured result"""
|
||
|
|
with await ssh_connection(config) as client:
|
||
|
|
stdin, stdout, stderr = await asyncio.to_thread(
|
||
|
|
client.exec_command, command, timeout=timeout
|
||
|
|
)
|
||
|
|
exit_code = stdout.channel.recv_exit_status()
|
||
|
|
out = await asyncio.to_thread(stdout.read)
|
||
|
|
err = await asyncio.to_thread(stderr.read)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"status": "success" if exit_code == 0 else "failed",
|
||
|
|
"stdout": out.decode()[:10000],
|
||
|
|
"stderr": err.decode()[:10000],
|
||
|
|
"exit_code": exit_code,
|
||
|
|
}
|