Files
Nexus/server/api/webssh.py
T
Nexus Deploy b866e944ab fix: Code Review P0/P1 batches 1-3 and single-operator risk acceptance
Apply sync/install/auth/schedule/retry/agent/settings fixes from full
code review; document accepted WS and Agent legacy risks for solo ops.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-04 15:57:49 +08:00

399 lines
16 KiB
Python

"""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
token_tv = payload.get("tv")
if token_tv is not None:
if int(token_tv) != int(admin.token_version or 0):
return None, None
else:
# Legacy WebSSH tokens without tv — fall back to updated_at grace
token_updated = payload.get("updated", 0)
if token_updated and admin.updated_at:
admin_ts = int(admin.updated_at.timestamp())
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=4401, 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,
})
# Color + CWD: ANSI/256/truecolor for ls, grep, vim, etc.
shell.stdin.write(
"export TERM=xterm-256color COLORTERM=truecolor "
"FORCE_COLOR=1 CLICOLOR_FORCE=1\n"
)
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}")