Your Name
8b10f9d64a
- Fix shell creation failure on stale pooled SSH connections:
force-close and retry with fresh connection
- Fix empty error message in logs: use {type(e).__name__}: {e!r}
- Fix double WebSocket.close() causing RuntimeError in finally block:
track ws_closed flag to avoid closing twice
- Fix Alpine undefined error in terminal.html: $d() now returns
default object when Alpine hasn't initialized yet
- All websocket.close() calls now wrapped in try/except
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
374 lines
14 KiB
Python
374 lines
14 KiB
Python
"""Nexus — Web SSH Terminal WebSocket Endpoint
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime
|
|
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"
|
|
|
|
# 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 to {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:
|
|
pass
|
|
try:
|
|
await websocket.close(code=4003, reason="SSH connection failed")
|
|
except Exception:
|
|
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),
|
|
)
|
|
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),
|
|
)
|
|
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:
|
|
pass
|
|
try:
|
|
await websocket.close(code=4003, reason="Shell creation failed")
|
|
except Exception:
|
|
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,
|
|
})
|
|
|
|
# ── Bidirectional relay ──
|
|
command_buffer = "" # Buffer for command logging
|
|
|
|
async def _read_shell_output():
|
|
"""Read SSH shell output -> send to WebSocket client"""
|
|
try:
|
|
async for data in shell.stdout:
|
|
await websocket.send_json({
|
|
"type": MSG_DATA,
|
|
"data": data,
|
|
})
|
|
except Exception:
|
|
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_CLOSE:
|
|
break
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception:
|
|
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:
|
|
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:
|
|
pass
|
|
|
|
try:
|
|
await websocket.close()
|
|
except Exception:
|
|
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 disconnected from {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}")
|