101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
import asyncio
|
||
import asyncssh
|
||
import os
|
||
import logging
|
||
from typing import Callable, Awaitable
|
||
from django.conf import settings
|
||
|
||
CHUNK_LIMIT = 32768 # store last N chars for tail
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class SSHError(Exception):
|
||
pass
|
||
|
||
async def open_connection(remote_host):
|
||
"""Open an SSH connection honoring host settings.
|
||
Adds diagnostics & graceful fallback if key file missing.
|
||
"""
|
||
if remote_host.strict_host_key_checking and settings.STRICT_HOST_KEY_CHECKING:
|
||
known_hosts = os.path.expanduser(settings.KNOWN_HOSTS_PATH)
|
||
else:
|
||
known_hosts = None # disables host key verification (dev only)
|
||
|
||
client_keys = None
|
||
agent_fwd = False
|
||
|
||
if remote_host.auth_method == 'ssh_key':
|
||
key_paths = []
|
||
if remote_host.key_path:
|
||
kp = os.path.expanduser(remote_host.key_path)
|
||
if os.path.exists(kp):
|
||
key_paths.append(kp)
|
||
logger.info("Using explicit SSH key: %s", kp)
|
||
else:
|
||
logger.warning("Configured key_path does not exist: %s (falling back to default keys / agent)", kp)
|
||
# If no valid explicit key, allow asyncssh to try default keys & agent
|
||
client_keys = key_paths or None
|
||
elif remote_host.auth_method == 'agent':
|
||
agent_fwd = True
|
||
client_keys = None # rely on agent
|
||
logger.info("Using SSH agent for auth (SSH_AUTH_SOCK=%s)", os.environ.get('SSH_AUTH_SOCK'))
|
||
elif remote_host.auth_method == 'password':
|
||
raise SSHError("Password auth not implemented – bitte SSH Key oder Agent nutzen")
|
||
|
||
try:
|
||
logger.info(
|
||
"Opening SSH connection to %s, port %s (known_hosts=%s, client_keys=%s, agent=%s)",
|
||
remote_host.hostname, remote_host.port, known_hosts if known_hosts else 'DISABLED', client_keys, agent_fwd
|
||
)
|
||
return await asyncssh.connect(
|
||
host=remote_host.hostname,
|
||
port=remote_host.port,
|
||
username=remote_host.username,
|
||
known_hosts=known_hosts,
|
||
client_keys=client_keys,
|
||
agent_forwarding=agent_fwd,
|
||
)
|
||
except asyncssh.HostKeyNotVerifiable as e:
|
||
raise SSHError(f"Host key not trusted for host {remote_host.hostname}") from e
|
||
except asyncssh.PermissionDenied as e:
|
||
raise SSHError(f"Permission denied for user {remote_host.username} on host {remote_host.hostname}") from e
|
||
except (asyncssh.DisconnectError, asyncssh.ConnectionLost) as e:
|
||
raise SSHError(f"Connection lost: {e}") from e
|
||
except (asyncssh.Error, OSError) as e:
|
||
raise SSHError(str(e)) from e
|
||
|
||
async def run_command(conn: asyncssh.SSHClientConnection, command: str, on_chunk: Callable[[str,str], Awaitable[None]], cancel_event: asyncio.Event) -> int:
|
||
try:
|
||
proc = await conn.create_process(command)
|
||
except (asyncssh.Error, OSError) as e:
|
||
raise SSHError(str(e)) from e
|
||
|
||
tail_buf = ''
|
||
|
||
async def reader(stream, stream_name):
|
||
nonlocal tail_buf
|
||
async for line in stream:
|
||
if cancel_event.is_set():
|
||
break
|
||
data = line.rstrip('\n')
|
||
tail_buf = (tail_buf + data + '\n')[-CHUNK_LIMIT:]
|
||
await on_chunk(stream_name, data + '\n')
|
||
|
||
stdout_task = asyncio.create_task(reader(proc.stdout, 'stdout'))
|
||
stderr_task = asyncio.create_task(reader(proc.stderr, 'stderr'))
|
||
|
||
while True:
|
||
if cancel_event.is_set():
|
||
proc.terminate()
|
||
try:
|
||
await asyncio.wait_for(proc.wait_closed(), timeout=5)
|
||
except asyncio.TimeoutError:
|
||
proc.kill()
|
||
return 130 # typical canceled code
|
||
if proc.exit_status is not None:
|
||
break
|
||
await asyncio.sleep(0.1)
|
||
|
||
await stdout_task
|
||
await stderr_task
|
||
return proc.exit_status
|