atc-agents/api/ssh_terminal.py

"""Interactive SSH terminal bridge over WebSocket (paramiko PTY <-> xterm.js)."""

from __future__ import annotations

import asyncio
import json
from typing import Any

import paramiko
from fastapi import WebSocket, WebSocketDisconnect


async def _send(ws: WebSocket, type_: str, **kw: Any) -> None:
    try:
        await ws.send_text(json.dumps({"type": type_, **kw}))
    except Exception:
        pass


def _open_ssh(cfg: dict[str, Any]) -> tuple[paramiko.SSHClient, paramiko.Channel]:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    connect_kwargs: dict[str, Any] = {
        "hostname": cfg["host"],
        "port": int(cfg.get("port") or 22),
        "username": cfg.get("username") or "root",
        "timeout": 12,
        "banner_timeout": 12,
        "auth_timeout": 12,
        "look_for_keys": False,
        "allow_agent": False,
    }
    password = cfg.get("password")
    key_data = cfg.get("private_key")
    if key_data:
        from io import StringIO
        pkey = None
        for loader in (paramiko.Ed25519Key, paramiko.RSAKey, paramiko.ECDSAKey):
            try:
                pkey = loader.from_private_key(StringIO(key_data), password=password or None)
                break
            except Exception:
                continue
        if pkey is None:
            raise ValueError("Unsupported or invalid private key")
        connect_kwargs["pkey"] = pkey
    elif password:
        connect_kwargs["password"] = password
    else:
        raise ValueError("No password or private key provided")

    client.connect(**connect_kwargs)
    cols = int(cfg.get("cols") or 120)
    rows = int(cfg.get("rows") or 32)
    chan = client.invoke_shell(term="xterm-256color", width=cols, height=rows)
    chan.settimeout(0.0)
    return client, chan


async def ssh_session(ws: WebSocket) -> None:
    await ws.accept()
    client: paramiko.SSHClient | None = None
    chan: paramiko.Channel | None = None
    loop = asyncio.get_event_loop()

    try:
        # First message must be the connect config
        first = await ws.receive_text()
        cfg = json.loads(first)
        if cfg.get("type") != "connect":
            await _send(ws, "error", message="Expected connect message")
            await ws.close()
            return

        await _send(ws, "status", message=f"Connecting to {cfg.get('username','root')}@{cfg.get('host')}:{cfg.get('port',22)}…")
        try:
            client, chan = await loop.run_in_executor(None, _open_ssh, cfg)
        except paramiko.AuthenticationException:
            await _send(ws, "error", message="Authentication failed — check username/password")
            await ws.close()
            return
        except Exception as exc:
            await _send(ws, "error", message=f"Connection failed: {exc}")
            await ws.close()
            return

        await _send(ws, "connected", message="connected")

        initial_cmd = cfg.get("initial_command")
        if initial_cmd and chan and not chan.closed:
            try:
                chan.send(initial_cmd if initial_cmd.endswith("\n") else initial_cmd + "\n")
            except Exception:
                pass

        async def pump_out() -> None:
            assert chan is not None
            while True:
                if chan.closed or chan.exit_status_ready() and not chan.recv_ready():
                    if not chan.recv_ready():
                        break
                if chan.recv_ready():
                    try:
                        data = await loop.run_in_executor(None, chan.recv, 65536)
                    except Exception:
                        break
                    if not data:
                        break
                    await ws.send_text(json.dumps({"type": "data", "data": data.decode("utf-8", errors="replace")}))
                else:
                    await asyncio.sleep(0.02)
            await _send(ws, "closed", message="Session closed")

        out_task = asyncio.create_task(pump_out())

        try:
            while True:
                msg = await ws.receive_text()
                try:
                    parsed = json.loads(msg)
                except json.JSONDecodeError:
                    parsed = {"type": "data", "data": msg}
                mtype = parsed.get("type")
                if mtype == "data" and chan and not chan.closed:
                    chan.send(parsed.get("data", ""))
                elif mtype == "resize" and chan and not chan.closed:
                    try:
                        chan.resize_pty(width=int(parsed.get("cols", 120)), height=int(parsed.get("rows", 32)))
                    except Exception:
                        pass
                elif mtype == "disconnect":
                    break
        finally:
            out_task.cancel()
    except WebSocketDisconnect:
        pass
    except Exception as exc:
        await _send(ws, "error", message=str(exc))
    finally:
        try:
            if chan is not None:
                chan.close()
        except Exception:
            pass
        try:
            if client is not None:
                client.close()
        except Exception:
            pass
        try:
            await ws.close()
        except Exception:
            pass