85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
from dataclasses import dataclass
|
|
|
|
import paramiko
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
class SSHClientError(Exception):
|
|
"""Raised when the SSH client cannot complete an operation."""
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class SSHConnectionConfig:
|
|
host: str
|
|
port: int
|
|
username: str
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class SSHCommandResult:
|
|
stdout: str
|
|
stderr: str
|
|
exit_code: int
|
|
|
|
|
|
class SSHClient:
|
|
def run_command(
|
|
self,
|
|
connection: SSHConnectionConfig,
|
|
command: str,
|
|
allowed_exit_codes: tuple[int, ...] = (0,),
|
|
) -> SSHCommandResult:
|
|
client = paramiko.SSHClient()
|
|
client.load_system_host_keys()
|
|
if settings.ssh_allow_unknown_host_keys:
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
connect_kwargs: dict[str, object] = {
|
|
"hostname": connection.host,
|
|
"port": connection.port,
|
|
"username": connection.username,
|
|
"timeout": settings.ssh_connect_timeout_seconds,
|
|
"banner_timeout": settings.ssh_connect_timeout_seconds,
|
|
"auth_timeout": settings.ssh_connect_timeout_seconds,
|
|
}
|
|
if settings.ssh_private_key_path:
|
|
connect_kwargs["key_filename"] = settings.ssh_private_key_path
|
|
|
|
try:
|
|
client.connect(**connect_kwargs)
|
|
stdin, stdout, stderr = client.exec_command(
|
|
command,
|
|
timeout=settings.ssh_command_timeout_seconds,
|
|
)
|
|
stdin.close()
|
|
exit_code = stdout.channel.recv_exit_status()
|
|
stdout_text = stdout.read().decode("utf-8").strip()
|
|
stderr_text = stderr.read().decode("utf-8").strip()
|
|
if exit_code not in allowed_exit_codes:
|
|
raise SSHClientError(
|
|
f"Command failed with exit code {exit_code}: {stderr_text or stdout_text}"
|
|
)
|
|
return SSHCommandResult(
|
|
stdout=stdout_text,
|
|
stderr=stderr_text,
|
|
exit_code=exit_code,
|
|
)
|
|
except Exception as exc: # pragma: no cover - library/transport failure path
|
|
if isinstance(exc, SSHClientError):
|
|
raise
|
|
raise SSHClientError(str(exc)) from exc
|
|
finally:
|
|
client.close()
|
|
|
|
def run_commands(
|
|
self,
|
|
connection: SSHConnectionConfig,
|
|
commands: dict[str, str],
|
|
) -> dict[str, str]:
|
|
output: dict[str, str] = {}
|
|
for key, command in commands.items():
|
|
result = self.run_command(connection, command)
|
|
output[key] = result.stdout
|
|
return output
|