"""Nexus — Web SSH Terminal WebSocket Endpoint """ import asyncio import json import logging import uuid from typing import Optional from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query, Path from server.config import settings from server.domain.models import SshSession, CommandLog, Server, AuditLog from server.infrastructure.ssh.asyncssh_pool import ssh_pool from server.infrastructure.database.session import AsyncSessionLocal from server.infrastructure.database.admin_repo import AdminRepositoryImpl from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.ssh_session_repo import SshSessionRepositoryImpl, CommandLogRepositoryImpl from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl logger = logging.getLogger("nexus.webssh") router = APIRouter() # Koko protocol message types MSG_TERMINAL_INIT = "TERMINAL_INIT" MSG_DATA = "DATA" MSG_RESIZE = "RESIZE" MSG_CLOSE = "CLOSE" MSG_ERROR = "ERROR" MSG_PING = "PING" MSG_PONG = "PONG" # Command logging configuration COMMAND_LOG_MAX_LENGTH = 2000 async def _verify_webssh_token(token: str): """Verify JWT token for WebSSH — returns (admin, server_id) or (None, None) Checks: valid signature, not expired, sub exists, admin active, and updated_at matches (invalidates token after password change). """ try: import jwt as pyjwt payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"], options={"require": ["exp", "sub"]}) if payload.get("purpose") != "webssh": return None, None admin_id = payload.get("sub") server_id = payload.get("server_id") if not admin_id or server_id is None: return None, None try: server_id = int(server_id) except (TypeError, ValueError): return None, None async with AsyncSessionLocal() as session: admin_repo = AdminRepositoryImpl(session) admin = await admin_repo.get_by_id(admin_id) if not admin or not admin.is_active: return None, None # P1-7: Check updated_at — invalidate token after password change token_updated = payload.get("updated", 0) if token_updated and admin.updated_at: admin_ts = int(admin.updated_at.timestamp()) # Allow 5-second grace for clock skew / race conditions if admin_ts > token_updated + 5: return None, None return admin, server_id except Exception: logger.debug("WebSSH JWT verification failed", exc_info=True) return None, None @router.websocket("/ws/terminal/{server_id}") async def terminal_ws( websocket: WebSocket, server_id: int = Path(..., description="Target server ID"), token: Optional[str] = Query(None, description="JWT access token"), ): """WebSSH Terminal WebSocket endpoint ADR-011: JWT authentication via query parameter `?token=xxx` W3: WebSocket endpoint for terminal access W5: Koko message protocol """ # ── JWT Authentication ── if not token: await websocket.close(code=4001, reason="Missing JWT token") return admin, token_server_id = await _verify_webssh_token(token) if not admin: await websocket.close(code=4001, reason="Invalid or expired JWT token") return # Security: WebSSH token must bind to this server (general access_token rejected) if token_server_id != server_id: await websocket.close(code=4003, reason="Token not authorized for this server") return # ── P2-16: Consolidated pre-accept DB operations in one session ── # server lookup + audit log + SSH session record all share one session, # which is closed before the long-lived WebSocket relay begins. client_ip = websocket.client.host if websocket.client else "unknown" session_id = str(uuid.uuid4()) server: Optional[Server] = None async with AsyncSessionLocal() as db: server_repo = ServerRepositoryImpl(db) server = await server_repo.get_by_id(server_id) if not server: await websocket.close(code=4003, reason="Server not found") return # ── Server-level authorization check ── if not server.domain: await websocket.close(code=4003, reason="Server domain not configured") return if server.auth_method == "key" and not server.ssh_key_path and not server.ssh_key_private: await websocket.close(code=4003, reason="SSH key not configured") return if server.auth_method == "password" and not server.password: await websocket.close(code=4003, reason="SSH password not configured") return # Audit: WebSSH session start audit_repo = AuditLogRepositoryImpl(db) await audit_repo.create(AuditLog( admin_username=admin.username, action="webssh_connect", target_type="server", target_id=server.id, detail=f"WebSSH 连接 {server.name} ({server.domain})", ip_address=client_ip, )) # Create SSH session record ssh_session_repo = SshSessionRepositoryImpl(db) ssh_session = SshSession( id=session_id, server_id=server.id, admin_id=admin.id, remote_addr=client_ip, status="active", ) await ssh_session_repo.create(ssh_session) # ── Accept WebSocket ── await websocket.accept() logger.info(f"WebSSH: admin={admin.username} connecting to server={server.name} ({server.domain})") # ── Establish SSH connection (with stale-connection retry) ── conn = None try: conn = await ssh_pool.acquire(server) except Exception as e: logger.error(f"WebSSH connection failed: {e}") try: await websocket.send_json({"type": MSG_ERROR, "message": f"SSH连接失败: {str(e)}"}) except Exception: # noqa: S110 — best-effort cleanup pass try: await websocket.close(code=4003, reason="SSH connection failed") except Exception: # noqa: S110 — best-effort cleanup pass await _close_ssh_session(session_id) return # ── Create SSH shell (retry once if stale pooled connection) ── ws_closed = False # Track whether we've closed the WebSocket try: try: shell_proc = conn.create_process( term_type="xterm-256color", term_size=(24, 80), bufsize=4096, # Small buffer for low-latency interactive echo ) except Exception as e: # Stale pooled connection — force-close and retry with fresh connection logger.warning(f"WebSSH shell creation failed on pooled connection (server={server.id}): {type(e).__name__}: {e!r}") await ssh_pool.close_connection(server.id) try: conn = await ssh_pool.acquire(server) shell_proc = conn.create_process( term_type="xterm-256color", term_size=(24, 80), bufsize=4096, # Small buffer for low-latency interactive echo ) except Exception as e2: logger.error(f"WebSSH shell creation failed on fresh connection (server={server.id}): {type(e2).__name__}: {e2!r}") try: await websocket.send_json({"type": MSG_ERROR, "message": f"Shell创建失败: {type(e2).__name__}"}) except Exception: # noqa: S110 — best-effort cleanup pass try: await websocket.close(code=4003, reason="Shell creation failed") except Exception: # noqa: S110 — best-effort cleanup pass ws_closed = True await _close_ssh_session(session_id) return async with shell_proc as shell: # Send TERMINAL_INIT to client await websocket.send_json({ "type": MSG_TERMINAL_INIT, "session_id": session_id, "server_name": server.name, "server_id": server.id, }) # Inject CWD tracking via OSC title escape shell.stdin.write('export PROMPT_COMMAND=\'echo -ne "\\033]0;${PWD}\\a"\'\n') # ── Bidirectional relay ── command_buffer = "" # Buffer for command logging async def _read_shell_output(): """Read SSH shell output -> send to WebSocket client Uses read() instead of async-for to avoid readline() buffering. async-for calls readline() internally, which waits for newline characters before yielding — causing echo delay in interactive terminal use. read(n) returns data as soon as any is available. """ try: while not shell.stdout.at_eof(): data = await shell.stdout.read(4096) if data: await websocket.send_json({ "type": MSG_DATA, "data": data, }) except Exception: # noqa: S110 — SSH read loop, connection teardown pass async def _read_websocket_input(): """Read WebSocket client input -> write to SSH shell""" nonlocal command_buffer try: while True: raw = await websocket.receive_text() try: msg = json.loads(raw) except json.JSONDecodeError: # Raw text input (backward compat) shell.stdin.write(raw) command_buffer += raw continue msg_type = msg.get("type", "") if msg_type == MSG_DATA: data = msg.get("data", "") shell.stdin.write(data) # Buffer for command logging command_buffer += data # Log complete commands (on Enter) if "\r" in data or "\n" in data: cmd = command_buffer.strip() if cmd: # P2-18: Check for dangerous commands from server.api.dependencies import check_dangerous_command check_dangerous_command(cmd) await _log_command( session_id=session_id, server_id=server.id, admin_id=admin.id, command=cmd[:COMMAND_LOG_MAX_LENGTH], remote_addr=client_ip, ) command_buffer = "" elif msg_type == MSG_RESIZE: cols = msg.get("cols", 80) rows = msg.get("rows", 24) shell.change_terminal_size(rows, cols) elif msg_type == MSG_PING: await websocket.send_json({"type": MSG_PONG, "ts": msg.get("ts")}) elif msg_type == MSG_CLOSE: break except WebSocketDisconnect: pass except Exception: # noqa: S110 — WebSocket read loop, connection teardown pass # Run both directions concurrently read_task = asyncio.create_task(_read_shell_output()) write_task = asyncio.create_task(_read_websocket_input()) try: # Wait for either direction to finish done, pending = await asyncio.wait( [read_task, write_task], return_when=asyncio.FIRST_COMPLETED, ) for task in pending: task.cancel() try: await task except asyncio.CancelledError: pass except Exception as e: logger.warning(f"WebSSH relay error: {e}") except Exception as e: logger.error(f"WebSSH shell creation failed: {type(e).__name__}: {e!r}") if not ws_closed: try: await websocket.send_json({"type": MSG_ERROR, "message": f"Shell创建失败: {type(e).__name__}"}) except Exception: # noqa: S110 — best-effort cleanup pass finally: # ── Cleanup ── await ssh_pool.release(server.id) await _close_ssh_session(session_id) if not ws_closed: try: await websocket.send_json({"type": MSG_CLOSE, "session_id": session_id}) except Exception: # noqa: S110 — best-effort cleanup pass try: await websocket.close() except Exception: # noqa: S110 — best-effort cleanup pass logger.info(f"WebSSH session closed: admin={admin.username}, server={server.name}, session={session_id}") # Audit: WebSSH session end try: async with AsyncSessionLocal() as audit_session: audit_repo = AuditLogRepositoryImpl(audit_session) await audit_repo.create(AuditLog( admin_username=admin.username, action="webssh_disconnect", target_type="server", target_id=server.id, detail=f"WebSSH 断开 {server.name}", ip_address=client_ip, )) except Exception: logger.warning(f"Failed to write disconnect audit log for session {session_id}") async def _close_ssh_session(session_id: str): """Close SSH session record in database""" try: async with AsyncSessionLocal() as session: ssh_session_repo = SshSessionRepositoryImpl(session) await ssh_session_repo.close_session(session_id) except Exception as e: logger.error(f"Failed to close SSH session {session_id}: {e}") async def _log_command(session_id: str, server_id: int, admin_id: int, command: str, remote_addr: str): """Log a command to CommandLog Uses its own session per command — the WebSocket connection is long-lived (potentially hours), so we can't hold a single DB session open for the entire duration. Each command gets its own short-lived session. """ try: async with AsyncSessionLocal() as session: cmd_log_repo = CommandLogRepositoryImpl(session) log = CommandLog( session_id=session_id, server_id=server_id, admin_id=admin_id, command=command, remote_addr=remote_addr, ) await cmd_log_repo.create(log) except Exception as e: logger.error(f"Failed to log command: {e}")