nas-tools/dashboard/backend/routers/terminal.py

import asyncio
import logging
import shlex
import struct
from collections import Counter

import asyncssh
import jwt
from fastapi import Query, WebSocket

import auth_service as auth
import config

logger = logging.getLogger(__name__)

MAX_SESSIONS = 5
MAX_SESSIONS_PER_USER = 5
_session_lock = asyncio.Lock()
_active_sessions: dict[int, str] = {}
PERSISTENCE_STATUS_PREFIX = "__DASHBOARD_PERSISTENCE__:"

HOSTS = {
    "nas": {"host": config.SSH_HOST, "user": config.SSH_USER, "key": config.SSH_KEY_PATH},
    "vps": {"host": config.VPS_SSH_HOST, "user": config.VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
    "caddy-vps": {"host": config.CADDY_VPS_SSH_HOST, "user": config.CADDY_VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
    "mac": {"host": config.MAC_SSH_HOST, "user": config.MAC_SSH_USER, "key": config.MAC_SSH_KEY_PATH},
    "claude-dev": {
        "host": config.SSH_HOST,
        "user": config.SSH_USER,
        "key": config.SSH_KEY_PATH,
        "persistent_session": {
            "container": "claude-dev",
            "session_name": "dashboard-terminal",
        },
    },
}


# Load known_hosts once at import time — fail closed if missing
_known_hosts_available = False
try:
    _known_hosts = asyncssh.read_known_hosts(config.SSH_KNOWN_HOSTS)
    _known_hosts_available = True
except Exception:
    _known_hosts = None
    logger.error("SSH known_hosts file not found at %s — terminal connections will be refused", config.SSH_KNOWN_HOSTS)


def build_host_command(host_config: dict) -> str | None:
    persistent = host_config.get("persistent_session")
    if not persistent:
        return host_config.get("command")

    docker_bin = "/volume1/@appstore/ContainerManager/usr/bin/docker"
    container = persistent["container"]
    session_name = persistent["session_name"]
    quoted_container = shlex.quote(container)
    quoted_session = shlex.quote(session_name)
    quoted_marker = shlex.quote(PERSISTENCE_STATUS_PREFIX)
    # Use login shell to ensure env vars are loaded from .bashrc
    attach_cmd = (
        f"{docker_bin} exec -it {quoted_container} "
        f"bash -l -c 'cd /repos/nas-tools && tmux new-session -A -s {quoted_session}'"
    )

    script = " ".join(
        [
            "if",
            f"{docker_bin} exec {quoted_container} tmux has-session -t {quoted_session} >/dev/null 2>&1;",
            "then",
            f"printf '%sattached\\n' {quoted_marker};",
            "else",
            f"printf '%screated\\n' {quoted_marker};",
            "fi;",
            f"exec {attach_cmd}",
        ]
    )
    return f"bash -lc {shlex.quote(script)}"


async def _reserve_session(session_id: int, username: str) -> tuple[bool, str | None, int | None]:
    async with _session_lock:
        if len(_active_sessions) >= MAX_SESSIONS:
            return False, "Too many sessions", 1013
        user_sessions = Counter(_active_sessions.values())
        if user_sessions[username] >= MAX_SESSIONS_PER_USER:
            return False, "Too many sessions for user", 1013
        _active_sessions[session_id] = username
        return True, None, session_id


async def _release_session(session_id: int | None):
    if session_id is None:
        return
    async with _session_lock:
        _active_sessions.pop(session_id, None)


async def ws_endpoint(websocket: WebSocket, host: str = Query("nas")):
    if host not in HOSTS:
        await websocket.close(code=1008, reason="Unknown host")
        return

    access_token = websocket.cookies.get(auth.ACCESS_COOKIE_NAME)
    if not access_token:
        logger.warning("WebSocket auth failed for %s: no access token cookie", host)
        await websocket.close(code=1008, reason="Unauthorized")
        return

    try:
        user = await auth.get_current_user_ws(access_token)
    except jwt.ExpiredSignatureError as e:
        logger.warning("WebSocket auth failed for %s: token expired - %s", host, e)
        await websocket.close(code=1008, reason="Token expired")
        return
    except jwt.InvalidTokenError as e:
        logger.warning("WebSocket auth failed for %s: invalid token - %s", host, e)
        await websocket.close(code=1008, reason="Invalid token")
        return
    except Exception as e:
        logger.error("WebSocket auth failed for %s: unexpected error - %s", host, e)
        await websocket.close(code=1008, reason="Unauthorized")
        return

    if user.pages != "*" and "terminal" not in user.pages:
        await websocket.close(code=1008, reason="Access denied")
        return

    if not _known_hosts_available:
        await websocket.close(code=1011, reason="SSH host verification unavailable")
        return

    session_id = id(websocket)
    reserved, reason, session_or_code = await _reserve_session(session_id, user.username)
    if not reserved:
        await websocket.close(code=session_or_code, reason=reason)
        return
    session_id = session_or_code

    await websocket.accept()
    logger.info("WebSocket accepted for %s (user: %s)", host, user.username)

    h = HOSTS[host]
    try:
        conn = await asyncio.wait_for(
            asyncssh.connect(
                h["host"],
                username=h["user"],
                client_keys=[h["key"]],
                known_hosts=_known_hosts,
                keepalive_interval=15,
                keepalive_count_max=8,
            ),
            timeout=10.0,
        )
        logger.info("SSH connection established to %s", host)
    except TimeoutError:
        logger.error("SSH connection to %s timed out after 10s", host)
        await _release_session(session_id)
        await websocket.close(code=1011, reason="SSH connection timeout")
        return
    except Exception as e:
        logger.error("SSH connection to %s failed: %s", host, e)
        await _release_session(session_id)
        await websocket.close(code=1011, reason="SSH connection failed")
        return

    try:
        process = await conn.create_process(
            build_host_command(h),
            term_type="xterm-256color",
            term_size=(80, 24),
            encoding=None,
        )

        async def read_ssh():
            try:
                while True:
                    data = await process.stdout.read(4096)
                    if not data:
                        break
                    await websocket.send_bytes(data)
            except asyncssh.Error as e:
                logger.warning("SSH read error for %s: %s", host, e)
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error("Unexpected error reading SSH for %s: %s", host, e)

        read_task = asyncio.create_task(read_ssh())

        try:
            while True:
                msg = await websocket.receive()
                if "bytes" in msg and msg["bytes"]:
                    data = msg["bytes"]
                    if data[0:1] == b"\x01" and len(data) == 5:
                        cols = struct.unpack("!H", data[1:3])[0]
                        rows = struct.unpack("!H", data[3:5])[0]
                        process.channel.change_terminal_size(cols, rows)
                    else:
                        try:
                            process.stdin.write(data)
                        except (BrokenPipeError, OSError) as e:
                            logger.warning("Failed to write to stdin for %s: %s", host, e)
                            break
                elif "text" in msg and msg["text"]:
                    if msg["text"] == "__ping__":
                        logger.debug("Received ping from %s, sending pong", host)
                        await websocket.send_text("__pong__")
                        continue
                    try:
                        process.stdin.write(msg["text"].encode())
                    except (BrokenPipeError, OSError) as e:
                        logger.warning("Failed to write to stdin for %s: %s", host, e)
                        break
        except Exception as e:
            logger.warning("WebSocket receive loop ended for %s: %s", host, e)
        finally:
            read_task.cancel()
            try:
                await asyncio.wait_for(read_task, timeout=2.0)
            except (TimeoutError, asyncio.CancelledError):
                pass
            process.close()
    finally:
        conn.close()
        await conn.wait_closed()
        await _release_session(session_id)