import asyncio import struct import logging import shlex import jwt from collections import Counter from fastapi import WebSocket, Query import asyncssh import config import auth_service as auth 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 asyncio.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 (asyncio.TimeoutError, asyncio.CancelledError): pass process.close() finally: conn.close() await conn.wait_closed() await _release_session(session_id)