488 lines
18 KiB
Python
488 lines
18 KiB
Python
"""
|
|
SSHService — asyncssh wrapper with per-device/per-object/global concurrency limits.
|
|
|
|
Each call to run_command() opens a fresh connection (no pooling).
|
|
Semaphores prevent overwhelming devices or the central server.
|
|
|
|
Auth resolution order for a device:
|
|
1. device.ssh_last_auth (last successful method — tried first)
|
|
2. device.ssh_key_id_override (specific SSH key for this device)
|
|
3. device.ssh_pass_override_enc (legacy password override for this device)
|
|
4. object_category_credentials WHERE category=device.category AND protocol='ssh'
|
|
5. object.ssh_credentials (ordered credential list)
|
|
6. object.ssh_user + object.ssh_pass_enc (legacy fallback)
|
|
|
|
On success the working method is persisted to device.ssh_last_auth.
|
|
If the last_auth method fails it is cleared and the full list is tried.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
import asyncssh
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.device import Device
|
|
from app.models.object import Object
|
|
from app.models.ssh_key import SSHKey
|
|
from app.services.crypto import decrypt
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Concurrency limits
|
|
_GLOBAL_LIMIT = 50
|
|
_OBJECT_LIMIT = 10
|
|
_DEVICE_LIMIT = 1
|
|
|
|
_global_sem = asyncio.Semaphore(_GLOBAL_LIMIT)
|
|
_object_sems: dict[int, asyncio.Semaphore] = {}
|
|
_device_sems: dict[int, asyncio.Semaphore] = {}
|
|
|
|
|
|
def _get_object_sem(object_id: int) -> asyncio.Semaphore:
|
|
if object_id not in _object_sems:
|
|
_object_sems[object_id] = asyncio.Semaphore(_OBJECT_LIMIT)
|
|
return _object_sems[object_id]
|
|
|
|
|
|
def _get_device_sem(device_id: int) -> asyncio.Semaphore:
|
|
if device_id not in _device_sems:
|
|
_device_sems[device_id] = asyncio.Semaphore(_DEVICE_LIMIT)
|
|
return _device_sems[device_id]
|
|
|
|
|
|
@dataclass
|
|
class SSHResult:
|
|
success: bool
|
|
stdout: str = ""
|
|
stderr: str = ""
|
|
exit_code: int = -1
|
|
duration_ms: int = 0
|
|
error: str = ""
|
|
# True when SSH auth succeeded (connection established), even if command returned non-zero.
|
|
# Used to distinguish "wrong password" from "command failed".
|
|
auth_ok: bool = False
|
|
|
|
|
|
@dataclass
|
|
class _AuthOption:
|
|
"""Internal representation of one auth attempt."""
|
|
type: str # "key" or "password"
|
|
username: str
|
|
password: str = "" # for type=password
|
|
key_path: str = "" # for type=key
|
|
# Identifiers used to persist last_auth
|
|
cred_id: Optional[int] = None # ObjectSSHCredential.id (if applicable)
|
|
key_id: Optional[int] = None # SSHKey.id (if applicable)
|
|
|
|
|
|
def _opt_key(opt: _AuthOption) -> str:
|
|
"""Deduplication key for an auth option."""
|
|
if opt.type == "key":
|
|
return f"key:{opt.key_id}:{opt.username}"
|
|
return f"password:{opt.cred_id}:{opt.username}"
|
|
|
|
|
|
async def _load_key_path(db: AsyncSession, key_id: int) -> Optional[str]:
|
|
"""Return the filesystem path for an SSHKey id, or None if not found."""
|
|
from sqlalchemy import select
|
|
result = await db.execute(select(SSHKey).where(SSHKey.id == key_id))
|
|
key = result.scalar_one_or_none()
|
|
return key.path if key else None
|
|
|
|
|
|
async def _add_global_credentials(
|
|
options: list[_AuthOption],
|
|
seen: set[str],
|
|
db: Optional[AsyncSession],
|
|
) -> None:
|
|
"""Append global SSH credentials (app-wide fallback) to the options list."""
|
|
if not db:
|
|
return
|
|
from sqlalchemy import select as sa_select
|
|
from app.models.ssh_key import GlobalSSHCredential
|
|
|
|
def _add(opt: _AuthOption) -> None:
|
|
k = _opt_key(opt)
|
|
if k not in seen:
|
|
seen.add(k)
|
|
options.append(opt)
|
|
|
|
result = await db.execute(
|
|
sa_select(GlobalSSHCredential).order_by(GlobalSSHCredential.priority, GlobalSSHCredential.id)
|
|
)
|
|
for gcred in result.scalars().all():
|
|
if gcred.type == "key" and gcred.ssh_key_id:
|
|
path = await _load_key_path(db, gcred.ssh_key_id)
|
|
if path:
|
|
_add(_AuthOption(type="key", username=gcred.username, key_path=path, key_id=gcred.ssh_key_id))
|
|
elif gcred.type == "password" and gcred.password_enc:
|
|
_add(_AuthOption(type="password", username=gcred.username, password=decrypt(gcred.password_enc)))
|
|
|
|
|
|
async def build_auth_options(device: Device, obj: Object, db: Optional[AsyncSession] = None) -> list[_AuthOption]:
|
|
"""Build an ordered list of auth options to try for this device.
|
|
|
|
If db is provided, key paths are resolved from the database.
|
|
Without db, key-based options referencing the DB are skipped.
|
|
"""
|
|
options: list[_AuthOption] = []
|
|
seen: set[str] = set()
|
|
|
|
def _add(opt: _AuthOption) -> None:
|
|
k = _opt_key(opt)
|
|
if k not in seen:
|
|
seen.add(k)
|
|
options.append(opt)
|
|
|
|
port = device.ssh_port_override or obj.ssh_port or 22 # not used in options, used by caller
|
|
|
|
# 1. Last successful method — try first
|
|
if device.ssh_last_auth:
|
|
la = device.ssh_last_auth
|
|
if la.get("type") == "key" and la.get("key_id") and db:
|
|
path = await _load_key_path(db, la["key_id"])
|
|
if path:
|
|
_add(_AuthOption(
|
|
type="key",
|
|
username=la.get("username", obj.ssh_user or "root"),
|
|
key_path=path,
|
|
key_id=la["key_id"],
|
|
))
|
|
elif la.get("type") == "password":
|
|
# We store the password reference; reconstruct from device/object
|
|
username = la.get("username", obj.ssh_user or "root")
|
|
cred_id = la.get("cred_id")
|
|
# Try to get password from matching cred
|
|
password = ""
|
|
if cred_id:
|
|
for cred in obj.ssh_credentials:
|
|
if cred.id == cred_id and cred.password_enc:
|
|
password = decrypt(cred.password_enc)
|
|
break
|
|
if not password:
|
|
if device.ssh_pass_override_enc:
|
|
password = decrypt(device.ssh_pass_override_enc)
|
|
elif obj.ssh_pass_enc:
|
|
password = decrypt(obj.ssh_pass_enc)
|
|
_add(_AuthOption(type="password", username=username, password=password, cred_id=cred_id))
|
|
|
|
# 2. Device-specific SSH key override
|
|
if device.ssh_key_id_override and db:
|
|
path = await _load_key_path(db, device.ssh_key_id_override)
|
|
if path:
|
|
username = device.ssh_user_override or obj.ssh_user or "root"
|
|
_add(_AuthOption(type="key", username=username, key_path=path, key_id=device.ssh_key_id_override))
|
|
|
|
# 3. Device-specific password override (legacy)
|
|
if device.ssh_pass_override_enc:
|
|
username = device.ssh_user_override or obj.ssh_user or "root"
|
|
password = decrypt(device.ssh_pass_override_enc)
|
|
_add(_AuthOption(type="password", username=username, password=password))
|
|
|
|
# 4. Per-category SSH credentials from object
|
|
for cat_cred in obj.category_credentials:
|
|
if cat_cred.category != device.category or cat_cred.protocol != "ssh":
|
|
continue
|
|
if cat_cred.type == "key" and cat_cred.ssh_key_id and db:
|
|
path = await _load_key_path(db, cat_cred.ssh_key_id)
|
|
if path:
|
|
_add(_AuthOption(type="key", username=cat_cred.username, key_path=path, key_id=cat_cred.ssh_key_id))
|
|
elif cat_cred.type == "password" and cat_cred.password_enc:
|
|
_add(_AuthOption(
|
|
type="password",
|
|
username=cat_cred.username,
|
|
password=decrypt(cat_cred.password_enc),
|
|
cred_id=cat_cred.id,
|
|
))
|
|
|
|
# 5. Object-level SSH credential list
|
|
for cred in obj.ssh_credentials:
|
|
if cred.type == "key" and cred.ssh_key_id and db:
|
|
path = await _load_key_path(db, cred.ssh_key_id)
|
|
if path:
|
|
_add(_AuthOption(type="key", username=cred.username, key_path=path, key_id=cred.ssh_key_id))
|
|
elif cred.type == "password" and cred.password_enc:
|
|
_add(_AuthOption(
|
|
type="password",
|
|
username=cred.username,
|
|
password=decrypt(cred.password_enc),
|
|
cred_id=cred.id,
|
|
))
|
|
|
|
# 6. Legacy object ssh_user + ssh_pass_enc fallback
|
|
if obj.ssh_user and obj.ssh_pass_enc:
|
|
_add(_AuthOption(
|
|
type="password",
|
|
username=obj.ssh_user,
|
|
password=decrypt(obj.ssh_pass_enc),
|
|
))
|
|
|
|
# 7. Global app-wide credentials (last resort)
|
|
await _add_global_credentials(options, seen, db)
|
|
|
|
return options
|
|
|
|
|
|
async def _try_connect(
|
|
host: str,
|
|
port: int,
|
|
opt: _AuthOption,
|
|
command: str,
|
|
timeout: int,
|
|
) -> SSHResult:
|
|
"""Attempt a single SSH connection with the given auth option."""
|
|
start = time.monotonic()
|
|
connect_kwargs: dict = dict(
|
|
host=host,
|
|
port=port,
|
|
username=opt.username,
|
|
known_hosts=None,
|
|
connect_timeout=10,
|
|
)
|
|
if opt.type == "key":
|
|
connect_kwargs["client_keys"] = [opt.key_path]
|
|
connect_kwargs["passphrase"] = None
|
|
else:
|
|
connect_kwargs["password"] = opt.password
|
|
|
|
try:
|
|
conn = await asyncio.wait_for(asyncssh.connect(**connect_kwargs), timeout=15)
|
|
async with conn:
|
|
result = await asyncio.wait_for(conn.run(command, check=False), timeout=timeout)
|
|
duration = int((time.monotonic() - start) * 1000)
|
|
return SSHResult(
|
|
success=result.exit_status == 0,
|
|
stdout=(result.stdout or "").strip(),
|
|
stderr=(result.stderr or "").strip(),
|
|
exit_code=result.exit_status or 0,
|
|
duration_ms=duration,
|
|
auth_ok=True, # Connection was established — auth succeeded
|
|
)
|
|
except asyncio.TimeoutError:
|
|
return SSHResult(
|
|
success=False,
|
|
error="Connection timed out",
|
|
duration_ms=int((time.monotonic() - start) * 1000),
|
|
)
|
|
except (asyncssh.DisconnectError, asyncssh.PermissionDenied, asyncssh.ConnectionLost) as e:
|
|
return SSHResult(
|
|
success=False,
|
|
error=str(e),
|
|
duration_ms=int((time.monotonic() - start) * 1000),
|
|
)
|
|
except Exception as e:
|
|
return SSHResult(
|
|
success=False,
|
|
error=str(e),
|
|
duration_ms=int((time.monotonic() - start) * 1000),
|
|
)
|
|
|
|
|
|
async def _save_last_auth(device: Device, db: AsyncSession, opt: _AuthOption) -> None:
|
|
"""Persist the successful auth method to device.ssh_last_auth."""
|
|
payload: dict = {"type": opt.type, "username": opt.username}
|
|
if opt.type == "key":
|
|
payload["key_id"] = opt.key_id
|
|
else:
|
|
if opt.cred_id:
|
|
payload["cred_id"] = opt.cred_id
|
|
device.ssh_last_auth = payload
|
|
try:
|
|
await db.flush()
|
|
except Exception:
|
|
pass # non-critical
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Legacy helper kept for callers that don't have device/object context
|
|
def resolve_ssh_credentials(device: Device, obj: Object) -> tuple[str, str, int]:
|
|
"""Returns (username, password, port) — device overrides take priority.
|
|
|
|
DEPRECATED: use run_command() with db= for full multi-auth support.
|
|
"""
|
|
user = device.ssh_user_override or obj.ssh_user or "root"
|
|
port = device.ssh_port_override or obj.ssh_port or 22
|
|
|
|
if device.ssh_pass_override_enc:
|
|
password = decrypt(device.ssh_pass_override_enc)
|
|
elif obj.ssh_pass_enc:
|
|
password = decrypt(obj.ssh_pass_enc)
|
|
else:
|
|
password = ""
|
|
|
|
return user, password, port
|
|
|
|
|
|
async def run_command(
|
|
device: Device,
|
|
obj: Object,
|
|
command: str,
|
|
timeout: int = 30,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> SSHResult:
|
|
"""Run command on device, trying all configured auth methods in order.
|
|
|
|
If db is provided:
|
|
- SSH keys are resolved from the database.
|
|
- The successful auth method is persisted to device.ssh_last_auth.
|
|
|
|
If last_auth method fails, it is cleared and remaining methods are tried.
|
|
"""
|
|
port = device.ssh_port_override or obj.ssh_port or 22
|
|
|
|
async with _global_sem:
|
|
async with _get_object_sem(device.object_id):
|
|
async with _get_device_sem(device.id):
|
|
options = await build_auth_options(device, obj, db)
|
|
|
|
if not options:
|
|
# Fallback: no creds configured at all
|
|
user = device.ssh_user_override or obj.ssh_user or "root"
|
|
options = [_AuthOption(type="password", username=user, password="")]
|
|
|
|
last_auth_option = options[0] if device.ssh_last_auth else None
|
|
last_result: Optional[SSHResult] = None
|
|
|
|
for i, opt in enumerate(options):
|
|
if opt.type == "key":
|
|
logger.info(
|
|
"SSH attempt %d/%d for %s — key: user=%s key_path=%s",
|
|
i + 1, len(options), device.ip, opt.username, opt.key_path,
|
|
)
|
|
else:
|
|
masked = (opt.password[:1] + "***" + opt.password[-1:]) if len(opt.password) > 2 else "***"
|
|
logger.info(
|
|
"SSH attempt %d/%d for %s — password: user=%s password=%s",
|
|
i + 1, len(options), device.ip, opt.username, masked,
|
|
)
|
|
|
|
result = await _try_connect(device.ip, port, opt, command, timeout)
|
|
|
|
if result.auth_ok:
|
|
# Connection was established — auth succeeded.
|
|
# Return the result regardless of exit_code (command may have failed, not auth).
|
|
logger.info(
|
|
"SSH auth succeeded for %s with %s(%s)",
|
|
device.ip, opt.type, opt.username,
|
|
)
|
|
if db:
|
|
await _save_last_auth(device, db, opt)
|
|
return result
|
|
|
|
# Auth failed (wrong password / key / connection error)
|
|
# If last_auth method failed, clear it
|
|
if i == 0 and opt is last_auth_option and db:
|
|
device.ssh_last_auth = None
|
|
try:
|
|
await db.flush()
|
|
except Exception:
|
|
pass
|
|
|
|
last_result = result
|
|
logger.info(
|
|
"SSH auth failed for %s with %s(%s): %s",
|
|
device.ip, opt.type, opt.username, result.error,
|
|
)
|
|
|
|
# All methods exhausted
|
|
if last_result:
|
|
last_result.error = f"All auth methods exhausted. Last error: {last_result.error}"
|
|
return last_result
|
|
return SSHResult(success=False, error="No auth methods available")
|
|
|
|
|
|
async def run_command_direct(
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
command: str,
|
|
timeout: int = 30,
|
|
) -> SSHResult:
|
|
"""Connect directly with explicit credentials (for object-level hosts like rack PCs).
|
|
|
|
Used by device_discovery when connecting to a specific IP that isn't a Device record.
|
|
"""
|
|
opt = _AuthOption(type="password", username=user, password=password)
|
|
async with _global_sem:
|
|
return await _try_connect(host, port, opt, command, timeout)
|
|
|
|
|
|
async def run_command_direct_with_options(
|
|
host: str,
|
|
port: int,
|
|
options: list[_AuthOption],
|
|
command: str,
|
|
timeout: int = 30,
|
|
) -> SSHResult:
|
|
"""Try multiple auth options for a non-Device host (e.g. rack PC during discovery)."""
|
|
async with _global_sem:
|
|
last_result: Optional[SSHResult] = None
|
|
for i, opt in enumerate(options):
|
|
if opt.type == "key":
|
|
logger.info(
|
|
"SSH attempt %d/%d for %s — key: user=%s key_path=%s",
|
|
i + 1, len(options), host, opt.username, opt.key_path,
|
|
)
|
|
else:
|
|
masked = (opt.password[:1] + "***" + opt.password[-1:]) if len(opt.password) > 2 else "***"
|
|
logger.info(
|
|
"SSH attempt %d/%d for %s — password: user=%s password=%s",
|
|
i + 1, len(options), host, opt.username, masked,
|
|
)
|
|
result = await _try_connect(host, port, opt, command, timeout)
|
|
if result.auth_ok:
|
|
logger.info("SSH auth succeeded for %s with %s(%s)", host, opt.type, opt.username)
|
|
return result
|
|
logger.info("SSH auth failed for %s with %s(%s): %s", host, opt.type, opt.username, result.error)
|
|
last_result = result
|
|
if last_result:
|
|
last_result.error = f"All auth methods exhausted. Last error: {last_result.error}"
|
|
return last_result
|
|
return SSHResult(success=False, error="No auth methods available")
|
|
|
|
|
|
async def build_object_auth_options(obj: Object, db: Optional[AsyncSession] = None) -> list[_AuthOption]:
|
|
"""Build auth options for connecting to object-level hosts (server_ip, rack PCs).
|
|
|
|
Uses the object's credential list + legacy ssh_user/ssh_pass_enc.
|
|
No device-level overrides.
|
|
"""
|
|
options: list[_AuthOption] = []
|
|
seen: set[str] = set()
|
|
|
|
def _add(opt: _AuthOption) -> None:
|
|
k = _opt_key(opt)
|
|
if k not in seen:
|
|
seen.add(k)
|
|
options.append(opt)
|
|
|
|
# Object SSH credential list
|
|
for cred in obj.ssh_credentials:
|
|
if cred.type == "key" and cred.ssh_key_id and db:
|
|
path = await _load_key_path(db, cred.ssh_key_id)
|
|
if path:
|
|
_add(_AuthOption(type="key", username=cred.username, key_path=path, key_id=cred.ssh_key_id))
|
|
elif cred.type == "password" and cred.password_enc:
|
|
_add(_AuthOption(
|
|
type="password",
|
|
username=cred.username,
|
|
password=decrypt(cred.password_enc),
|
|
cred_id=cred.id,
|
|
))
|
|
|
|
# Legacy fallback
|
|
if obj.ssh_user and obj.ssh_pass_enc:
|
|
_add(_AuthOption(type="password", username=obj.ssh_user, password=decrypt(obj.ssh_pass_enc)))
|
|
|
|
# Global app-wide credentials (last resort)
|
|
await _add_global_credentials(options, seen, db)
|
|
|
|
return options
|