107 lines
3.6 KiB
Python
107 lines
3.6 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.
|
|
"""
|
|
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
import asyncssh
|
|
|
|
from app.models.device import Device
|
|
from app.models.object import Object
|
|
from app.services.crypto import decrypt
|
|
|
|
# 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 = ""
|
|
|
|
|
|
def resolve_ssh_credentials(device: Device, obj: Object) -> tuple[str, str, int]:
|
|
"""Returns (username, password, port) — device overrides take priority."""
|
|
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,
|
|
) -> SSHResult:
|
|
user, password, port = resolve_ssh_credentials(device, obj)
|
|
start = time.monotonic()
|
|
|
|
async with _global_sem:
|
|
async with _get_object_sem(device.object_id):
|
|
async with _get_device_sem(device.id):
|
|
try:
|
|
conn = await asyncio.wait_for(
|
|
asyncssh.connect(
|
|
device.ip,
|
|
port=port,
|
|
username=user,
|
|
password=password,
|
|
known_hosts=None,
|
|
connect_timeout=10,
|
|
),
|
|
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,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
return SSHResult(success=False, error="Connection timed out", duration_ms=int((time.monotonic() - start) * 1000))
|
|
except asyncssh.DisconnectError as e:
|
|
return SSHResult(success=False, error=f"SSH disconnect: {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))
|