truenas-burnin/app/terminal.py

"""
WebSocket → asyncssh PTY bridge for the live terminal drawer tab.

Protocol
--------
Client  → server:  binary  = raw terminal input bytes
                   text    = JSON control message, e.g. {"type":"resize","cols":80,"rows":24}
Server  → client:  binary  = raw terminal output bytes
"""

import asyncio
import json
import logging

import asyncssh
from fastapi import WebSocket, WebSocketDisconnect

log = logging.getLogger(__name__)


async def handle(ws: WebSocket) -> None:
    """Accept a WebSocket connection and bridge it to an SSH PTY."""
    await ws.accept()

    from app.config import settings  # late import — avoids circular at module level

    # ── Guard: SSH must be configured ──────────────────────────────────────
    if not settings.ssh_host:
        await _send(ws,
            b"\r\n\x1b[33mSSH not configured.\x1b[0m "
            b"Set SSH Host in \x1b[1mSettings \u2192 SSH\x1b[0m first.\r\n"
        )
        await ws.close(1008)
        return

    connect_kw: dict = dict(
        host=settings.ssh_host,
        port=settings.ssh_port,
        username=settings.ssh_user,
        known_hosts=None,
    )

    if settings.ssh_key.strip():
        try:
            connect_kw["client_keys"] = [asyncssh.import_private_key(settings.ssh_key)]
        except Exception as exc:
            await _send(ws, f"\r\n\x1b[31mBad SSH key: {exc}\x1b[0m\r\n".encode())
            await ws.close(1011)
            return
    elif settings.ssh_password:
        connect_kw["password"] = settings.ssh_password
    else:
        await _send(ws,
            b"\r\n\x1b[33mNo SSH credentials configured.\x1b[0m "
            b"Set a password or private key in Settings.\r\n"
        )
        await ws.close(1008)
        return

    await _send(ws,
        f"\r\n\x1b[36mConnecting to {settings.ssh_host}\u2026\x1b[0m\r\n".encode()
    )

    # ── Open SSH connection ─────────────────────────────────────────────────
    try:
        async with asyncssh.connect(**connect_kw) as conn:
            process = await conn.create_process(
                term_type="xterm-256color",
                term_size=(80, 24),
                encoding=None,   # raw bytes — xterm.js handles encoding
            )
            await _send(ws, b"\r\n\x1b[32mConnected\x1b[0m\r\n\r\n")

            stop = asyncio.Event()

            async def ssh_to_ws() -> None:
                try:
                    async for chunk in process.stdout:
                        await ws.send_bytes(chunk)
                except Exception:
                    pass
                finally:
                    stop.set()

            async def ws_to_ssh() -> None:
                try:
                    while not stop.is_set():
                        msg = await ws.receive()
                        if msg["type"] == "websocket.disconnect":
                            break
                        if msg.get("bytes"):
                            process.stdin.write(msg["bytes"])
                        elif msg.get("text"):
                            try:
                                ctrl = json.loads(msg["text"])
                                if ctrl.get("type") == "resize":
                                    process.change_terminal_size(
                                        int(ctrl["cols"]), int(ctrl["rows"])
                                    )
                            except Exception:
                                pass
                except WebSocketDisconnect:
                    pass
                except Exception:
                    pass
                finally:
                    stop.set()

            t1 = asyncio.create_task(ssh_to_ws())
            t2 = asyncio.create_task(ws_to_ssh())

            _done, pending = await asyncio.wait(
                [t1, t2], return_when=asyncio.FIRST_COMPLETED
            )
            for t in pending:
                t.cancel()
                try:
                    await t
                except asyncio.CancelledError:
                    pass

    except asyncssh.PermissionDenied:
        await _send(ws, b"\r\n\x1b[31mSSH permission denied.\x1b[0m\r\n")
    except asyncssh.DisconnectError as exc:
        await _send(ws, f"\r\n\x1b[31mSSH disconnected: {exc}\x1b[0m\r\n".encode())
    except OSError as exc:
        await _send(ws, f"\r\n\x1b[31mCannot reach {settings.ssh_host}: {exc}\x1b[0m\r\n".encode())
    except Exception as exc:
        log.exception("Terminal WebSocket error")
        await _send(ws, f"\r\n\x1b[31mError: {exc}\x1b[0m\r\n".encode())
    finally:
        try:
            await ws.close()
        except Exception:
            pass


async def _send(ws: WebSocket, data: bytes) -> None:
    """Best-effort send — silently swallow errors if the socket is already gone."""
    try:
        await ws.send_bytes(data)
    except Exception:
        pass