Nexus/server/api/webssh.py

"""Nexus — Web SSH Terminal WebSocket Endpoint

W3: WebSocket endpoint /ws/terminal/{token}
W5: Koko message protocol (TERMINAL_INIT/DATA/RESIZE/CLOSE)

Flow:
1. Client connects via WebSocket with JWT token in query param
2. Server verifies JWT, looks up server by ID
3. Server creates asyncssh connection to target server
4. Bidirectional data relay: client ↔ WebSocket ↔ asyncssh
5. Session tracked in SshSession + commands logged to CommandLog
"""

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
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

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
COMMAND_LOG_BUFFER_LINES = 1  # Log every line immediately


async def _verify_webssh_token(token: str):
    """Verify JWT token for WebSSH — returns (admin, server_id) or (None, None)"""
    try:
        import jwt as pyjwt
        payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"])
        admin_id = payload.get("sub")
        server_id = payload.get("server_id")  # Embedded in token for WebSSH
        if not admin_id:
            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
            return admin, server_id
    except Exception:
        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: verify the JWT's server_id matches the URL path server_id
    if token_server_id is not None and token_server_id != server_id:
        await websocket.close(code=4003, reason="Token not authorized for this server")
        return

    # ── Lookup target server ──
    async with AsyncSessionLocal() as session:
        server_repo = ServerRepositoryImpl(session)
        server = await server_repo.get_by_id(server_id)

    if not server:
        await websocket.close(code=4004, reason=f"Server {server_id} not found")
        return

    # ── Server-level authorization check ──
    # Verify the server has valid SSH credentials configured
    if not server.domain:
        await websocket.close(code=4003, reason=f"Server {server.name} has no domain/IP 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=f"Server {server.name} has no SSH key configured")
        return

    if server.auth_method == "password" and not server.password:
        await websocket.close(code=4003, reason=f"Server {server.name} has no SSH password configured")
        return

    # ── Audit: WebSSH session start ──
    async with AsyncSessionLocal() as session:
        from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
        from server.domain.models import AuditLog
        audit_repo = AuditLogRepositoryImpl(session)
        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=websocket.client.host if websocket.client else "",
        ))

    # ── Accept WebSocket ──
    await websocket.accept()
    client_ip = websocket.client.host if websocket.client else "unknown"
    session_id = str(uuid.uuid4())

    logger.info(f"WebSSH: admin={admin.username} connecting to server={server.name} ({server.domain})")

    # ── Create SSH session record ──
    async with AsyncSessionLocal() as session:
        ssh_session_repo = SshSessionRepositoryImpl(session)
        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)

    # ── Establish SSH connection ──
    try:
        conn = await ssh_pool.acquire(server)
    except Exception as e:
        logger.error(f"WebSSH connection failed: {e}")
        await websocket.send_json({"type": MSG_ERROR, "message": f"SSH连接失败: {str(e)}"})
        await websocket.close(code=4003, reason="SSH connection failed")
        await _close_ssh_session(session_id)
        return

    # ── Create SSH shell ──
    try:
        async with conn.create_process(
            term_type="xterm-256color",
            term_size=(24, 80),
        ) 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:
                                    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: {e}")
        try:
            await websocket.send_json({"type": MSG_ERROR, "message": f"Shell创建失败: {str(e)}"})
        except Exception:
            pass
    finally:
        # ── Cleanup ──
        await ssh_pool.release(server.id)
        await _close_ssh_session(session_id)

        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:
                from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
                from server.domain.models import AuditLog
                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=websocket.client.host if websocket.client else "",
                ))
        except Exception:
            pass


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"""
    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}")