320 lines
8.3 KiB
Python
320 lines
8.3 KiB
Python
"""Web TUI 服务器
|
||
|
||
提供 Web 终端界面和 WebSocket 终端通信
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import uuid
|
||
from collections.abc import AsyncGenerator
|
||
from contextlib import asynccontextmanager
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import HTMLResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
|
||
from minenasai.core import get_logger, get_settings, setup_logging
|
||
from minenasai.webtui.auth import get_auth_manager
|
||
from minenasai.webtui.ssh_manager import SSHSession, get_ssh_manager
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
# 静态文件目录
|
||
STATIC_DIR = Path(__file__).parent / "static"
|
||
|
||
|
||
class TerminalConnection:
|
||
"""终端连接"""
|
||
|
||
def __init__(
|
||
self,
|
||
websocket: WebSocket,
|
||
session_id: str,
|
||
user_id: str,
|
||
) -> None:
|
||
self.websocket = websocket
|
||
self.session_id = session_id
|
||
self.user_id = user_id
|
||
self.ssh_session: SSHSession | None = None
|
||
self.authenticated = False
|
||
|
||
async def send_json(self, data: dict[str, Any]) -> None:
|
||
"""发送 JSON 消息"""
|
||
try:
|
||
await self.websocket.send_json(data)
|
||
except Exception as e:
|
||
logger.error("发送消息失败", error=str(e))
|
||
|
||
async def send_output(self, data: bytes) -> None:
|
||
"""发送终端输出"""
|
||
try:
|
||
await self.websocket.send_json({
|
||
"type": "output",
|
||
"data": data.decode("utf-8", errors="replace"),
|
||
})
|
||
except Exception as e:
|
||
logger.error("发送输出失败", error=str(e))
|
||
|
||
|
||
class ConnectionManager:
|
||
"""WebSocket 连接管理器"""
|
||
|
||
def __init__(self) -> None:
|
||
self.connections: dict[str, TerminalConnection] = {}
|
||
|
||
async def connect(
|
||
self,
|
||
websocket: WebSocket,
|
||
session_id: str,
|
||
user_id: str,
|
||
) -> TerminalConnection:
|
||
"""接受新连接"""
|
||
await websocket.accept()
|
||
conn = TerminalConnection(websocket, session_id, user_id)
|
||
self.connections[session_id] = conn
|
||
logger.info("终端连接建立", session_id=session_id)
|
||
return conn
|
||
|
||
async def disconnect(self, session_id: str) -> None:
|
||
"""断开连接"""
|
||
conn = self.connections.pop(session_id, None)
|
||
if conn and conn.ssh_session:
|
||
ssh_manager = get_ssh_manager()
|
||
await ssh_manager.close_session(session_id)
|
||
logger.info("终端连接断开", session_id=session_id)
|
||
|
||
def get_connection(self, session_id: str) -> TerminalConnection | None:
|
||
"""获取连接"""
|
||
return self.connections.get(session_id)
|
||
|
||
|
||
manager = ConnectionManager()
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||
"""应用生命周期管理
|
||
|
||
Args:
|
||
_app: FastAPI 应用实例（lifespan 标准签名要求）
|
||
"""
|
||
settings = get_settings()
|
||
setup_logging(settings.logging)
|
||
logger.info("Web TUI 服务启动", port=settings.webtui.port)
|
||
yield
|
||
|
||
# 清理所有连接
|
||
ssh_manager = get_ssh_manager()
|
||
await ssh_manager.close_all()
|
||
logger.info("Web TUI 服务关闭")
|
||
|
||
|
||
app = FastAPI(
|
||
title="MineNASAI Web TUI",
|
||
description="Web 终端界面",
|
||
version="0.1.0",
|
||
lifespan=lifespan,
|
||
)
|
||
|
||
# CORS 配置
|
||
settings = get_settings()
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 静态文件
|
||
if STATIC_DIR.exists():
|
||
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
|
||
|
||
|
||
@app.get("/")
|
||
async def index() -> HTMLResponse:
|
||
"""首页"""
|
||
index_file = STATIC_DIR / "index.html"
|
||
if index_file.exists():
|
||
return HTMLResponse(content=index_file.read_text(encoding="utf-8"))
|
||
return HTMLResponse(content="<h1>MineNASAI Web TUI</h1>")
|
||
|
||
|
||
@app.get("/health")
|
||
async def health() -> dict[str, str]:
|
||
"""健康检查"""
|
||
return {"status": "healthy"}
|
||
|
||
|
||
@app.get("/api/stats")
|
||
async def stats() -> dict[str, Any]:
|
||
"""获取统计信息"""
|
||
ssh_manager = get_ssh_manager()
|
||
auth_manager = get_auth_manager()
|
||
|
||
return {
|
||
"connections": len(manager.connections),
|
||
"ssh": ssh_manager.get_stats(),
|
||
"auth": auth_manager.get_stats(),
|
||
}
|
||
|
||
|
||
@app.post("/api/token")
|
||
async def generate_token(
|
||
user_id: str = "anonymous",
|
||
expires_in: int = 3600,
|
||
) -> dict[str, Any]:
|
||
"""生成访问令牌（用于测试）"""
|
||
auth_manager = get_auth_manager()
|
||
token = auth_manager.generate_token(user_id, expires_in)
|
||
return {
|
||
"token": token,
|
||
"user_id": user_id,
|
||
"expires_in": expires_in,
|
||
}
|
||
|
||
|
||
@app.websocket("/ws/terminal")
|
||
async def terminal_websocket(websocket: WebSocket) -> None:
|
||
"""终端 WebSocket 端点"""
|
||
session_id = str(uuid.uuid4())
|
||
conn: TerminalConnection | None = None
|
||
|
||
try:
|
||
conn = await manager.connect(websocket, session_id, "")
|
||
|
||
while True:
|
||
data = await websocket.receive_json()
|
||
msg_type = data.get("type")
|
||
|
||
if msg_type == "auth":
|
||
await handle_auth(conn, data)
|
||
|
||
elif msg_type == "input":
|
||
await handle_input(conn, data)
|
||
|
||
elif msg_type == "resize":
|
||
await handle_resize(conn, data)
|
||
|
||
elif msg_type == "pong":
|
||
pass # 心跳响应
|
||
|
||
else:
|
||
await conn.send_json({
|
||
"type": "error",
|
||
"message": f"未知消息类型: {msg_type}",
|
||
})
|
||
|
||
except WebSocketDisconnect:
|
||
pass
|
||
except Exception as e:
|
||
logger.error("WebSocket 错误", error=str(e), session_id=session_id)
|
||
finally:
|
||
if conn:
|
||
await manager.disconnect(session_id)
|
||
|
||
|
||
async def handle_auth(conn: TerminalConnection, data: dict[str, Any]) -> None:
|
||
"""处理认证"""
|
||
token = data.get("token", "")
|
||
auth_manager = get_auth_manager()
|
||
|
||
# 允许匿名访问（开发模式）
|
||
if token == "anonymous":
|
||
conn.authenticated = True
|
||
conn.user_id = "anonymous"
|
||
await conn.send_json({
|
||
"type": "auth_ok",
|
||
"session_id": conn.session_id,
|
||
"user_id": conn.user_id,
|
||
})
|
||
|
||
# 创建 SSH 会话
|
||
await create_ssh_session(conn)
|
||
return
|
||
|
||
# 验证令牌
|
||
auth_token = auth_manager.verify_token(token)
|
||
if auth_token is None:
|
||
await conn.send_json({
|
||
"type": "auth_error",
|
||
"message": "无效的令牌",
|
||
})
|
||
return
|
||
|
||
conn.authenticated = True
|
||
conn.user_id = auth_token.user_id
|
||
|
||
await conn.send_json({
|
||
"type": "auth_ok",
|
||
"session_id": conn.session_id,
|
||
"user_id": conn.user_id,
|
||
})
|
||
|
||
# 创建 SSH 会话
|
||
await create_ssh_session(conn)
|
||
|
||
|
||
async def create_ssh_session(conn: TerminalConnection) -> None:
|
||
"""创建 SSH 会话"""
|
||
ssh_manager = get_ssh_manager()
|
||
session = await ssh_manager.create_session(conn.session_id)
|
||
|
||
if session is None:
|
||
await conn.send_json({
|
||
"type": "error",
|
||
"message": "SSH 连接失败",
|
||
})
|
||
return
|
||
|
||
conn.ssh_session = session
|
||
|
||
# 设置输出回调
|
||
def on_output(data: bytes) -> None:
|
||
asyncio.create_task(conn.send_output(data))
|
||
|
||
session.set_output_callback(on_output)
|
||
|
||
# 开始读取输出
|
||
await session.start_reading()
|
||
|
||
|
||
async def handle_input(conn: TerminalConnection, data: dict[str, Any]) -> None:
|
||
"""处理终端输入"""
|
||
if not conn.authenticated:
|
||
await conn.send_json({
|
||
"type": "error",
|
||
"message": "未认证",
|
||
})
|
||
return
|
||
|
||
if conn.ssh_session is None:
|
||
await conn.send_json({
|
||
"type": "error",
|
||
"message": "SSH 会话未建立",
|
||
})
|
||
return
|
||
|
||
input_data = data.get("data", "")
|
||
await conn.ssh_session.write(input_data)
|
||
|
||
|
||
async def handle_resize(conn: TerminalConnection, data: dict[str, Any]) -> None:
|
||
"""处理终端大小调整"""
|
||
if conn.ssh_session is None:
|
||
return
|
||
|
||
cols = data.get("cols", 80)
|
||
rows = data.get("rows", 24)
|
||
|
||
await conn.ssh_session.resize(cols, rows)
|
||
await conn.send_json({
|
||
"type": "resize_ok",
|
||
"cols": cols,
|
||
"rows": rows,
|
||
})
|