787 lines
24 KiB
Python
787 lines
24 KiB
Python
"""Web TUI 服务器
|
||
|
||
提供 Web 终端界面、WebUI 配置界面和 WebSocket 终端通信
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
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, JSONResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from pydantic import BaseModel
|
||
|
||
from minenasai.core import get_logger, get_settings, setup_logging
|
||
from minenasai.core.config import Settings, save_config
|
||
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"
|
||
WEBUI_DIR = STATIC_DIR / "webui"
|
||
|
||
# 服务启动时间（用于计算运行时长）
|
||
_start_time: float = 0.0
|
||
|
||
|
||
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 标准签名要求）
|
||
"""
|
||
global _start_time
|
||
_start_time = time.time()
|
||
|
||
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")
|
||
|
||
# WebUI 静态文件
|
||
if WEBUI_DIR.exists():
|
||
app.mount("/webui/static", StaticFiles(directory=str(WEBUI_DIR)), name="webui_static")
|
||
|
||
|
||
@app.get("/")
|
||
async def index() -> HTMLResponse:
|
||
"""首页 - 重定向到 WebUI"""
|
||
return HTMLResponse(content="""
|
||
<html>
|
||
<head><meta http-equiv="refresh" content="0; url=/webui"></head>
|
||
<body><a href="/webui">跳转到控制台</a></body>
|
||
</html>
|
||
""")
|
||
|
||
|
||
@app.get("/terminal")
|
||
async def terminal_page() -> 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("/webui")
|
||
async def webui_index() -> HTMLResponse:
|
||
"""WebUI 控制台首页"""
|
||
index_file = WEBUI_DIR / "index.html"
|
||
if index_file.exists():
|
||
return HTMLResponse(content=index_file.read_text(encoding="utf-8"))
|
||
return HTMLResponse(content="<h1>MineNASAI WebUI</h1><p>WebUI 文件未找到</p>")
|
||
|
||
|
||
@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()
|
||
settings = get_settings()
|
||
|
||
# 计算运行时长
|
||
uptime_seconds = time.time() - _start_time if _start_time > 0 else 0
|
||
hours = int(uptime_seconds // 3600)
|
||
minutes = int((uptime_seconds % 3600) // 60)
|
||
uptime = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m"
|
||
|
||
return {
|
||
"status": "running",
|
||
"connections": len(manager.connections),
|
||
"agents": len(settings.agents.items),
|
||
"tasks": 0, # TODO: 从调度器获取
|
||
"uptime": uptime,
|
||
"ssh": ssh_manager.get_stats(),
|
||
"auth": auth_manager.get_stats(),
|
||
}
|
||
|
||
|
||
# ==================== 配置管理 API ====================
|
||
|
||
|
||
class LLMConfigUpdate(BaseModel):
|
||
"""LLM 配置更新"""
|
||
default_provider: str | None = None
|
||
default_model: str | None = None
|
||
anthropic_api_key: str | None = None
|
||
openai_api_key: str | None = None
|
||
deepseek_api_key: str | None = None
|
||
zhipu_api_key: str | None = None
|
||
minimax_api_key: str | None = None
|
||
minimax_group_id: str | None = None
|
||
moonshot_api_key: str | None = None
|
||
gemini_api_key: str | None = None
|
||
|
||
|
||
class ChannelConfigUpdate(BaseModel):
|
||
"""通讯渠道配置更新"""
|
||
wework: dict[str, Any] | None = None
|
||
feishu: dict[str, Any] | None = None
|
||
|
||
|
||
class ProxyConfigUpdate(BaseModel):
|
||
"""代理配置更新"""
|
||
enabled: bool | None = None
|
||
http: str | None = None
|
||
https: str | None = None
|
||
no_proxy: list[str] | None = None
|
||
auto_detect: bool | None = None
|
||
|
||
|
||
@app.get("/api/config")
|
||
async def get_config() -> dict[str, Any]:
|
||
"""获取配置（隐藏敏感信息）"""
|
||
settings = get_settings()
|
||
|
||
# 构建安全的配置响应
|
||
return {
|
||
"llm": {
|
||
"default_provider": settings.llm.default_provider,
|
||
"default_model": settings.llm.default_model,
|
||
# API Key 只显示是否已配置
|
||
"anthropic_api_key": "***" if settings.llm.anthropic_api_key else "",
|
||
"openai_api_key": "***" if settings.llm.openai_api_key else "",
|
||
"deepseek_api_key": "***" if settings.llm.deepseek_api_key else "",
|
||
"zhipu_api_key": "***" if settings.llm.zhipu_api_key else "",
|
||
"minimax_api_key": "***" if settings.llm.minimax_api_key else "",
|
||
"minimax_group_id": settings.llm.minimax_group_id or "",
|
||
"moonshot_api_key": "***" if settings.llm.moonshot_api_key else "",
|
||
"gemini_api_key": "***" if settings.llm.gemini_api_key else "",
|
||
},
|
||
"channels": {
|
||
"wework": {
|
||
"enabled": settings.channels.wework.enabled,
|
||
"corp_id": settings.wework_corp_id or "",
|
||
"agent_id": settings.wework_agent_id or "",
|
||
"secret": "***" if settings.wework_secret else "",
|
||
"token": settings.wework_token or "",
|
||
"encoding_aes_key": "***" if settings.wework_encoding_aes_key else "",
|
||
},
|
||
"feishu": {
|
||
"enabled": settings.channels.feishu.enabled,
|
||
"app_id": settings.feishu_app_id or "",
|
||
"app_secret": "***" if settings.feishu_app_secret else "",
|
||
"verification_token": settings.feishu_verification_token or "",
|
||
"encrypt_key": "***" if settings.feishu_encrypt_key else "",
|
||
}
|
||
},
|
||
"proxy": {
|
||
"enabled": settings.proxy.enabled,
|
||
"http": settings.proxy.http,
|
||
"https": settings.proxy.https,
|
||
"no_proxy": settings.proxy.no_proxy,
|
||
"auto_detect": settings.proxy.auto_detect,
|
||
}
|
||
}
|
||
|
||
|
||
@app.put("/api/config/llm")
|
||
async def update_llm_config(data: LLMConfigUpdate) -> dict[str, Any]:
|
||
"""更新 LLM 配置"""
|
||
settings = get_settings()
|
||
|
||
# 更新非空字段
|
||
if data.default_provider is not None:
|
||
settings.llm.default_provider = data.default_provider
|
||
if data.default_model is not None:
|
||
settings.llm.default_model = data.default_model
|
||
if data.anthropic_api_key and data.anthropic_api_key != "***":
|
||
settings.llm.anthropic_api_key = data.anthropic_api_key
|
||
if data.openai_api_key and data.openai_api_key != "***":
|
||
settings.llm.openai_api_key = data.openai_api_key
|
||
if data.deepseek_api_key and data.deepseek_api_key != "***":
|
||
settings.llm.deepseek_api_key = data.deepseek_api_key
|
||
if data.zhipu_api_key and data.zhipu_api_key != "***":
|
||
settings.llm.zhipu_api_key = data.zhipu_api_key
|
||
if data.minimax_api_key and data.minimax_api_key != "***":
|
||
settings.llm.minimax_api_key = data.minimax_api_key
|
||
if data.minimax_group_id:
|
||
settings.llm.minimax_group_id = data.minimax_group_id
|
||
if data.moonshot_api_key and data.moonshot_api_key != "***":
|
||
settings.llm.moonshot_api_key = data.moonshot_api_key
|
||
if data.gemini_api_key and data.gemini_api_key != "***":
|
||
settings.llm.gemini_api_key = data.gemini_api_key
|
||
|
||
# 保存配置
|
||
save_config(settings)
|
||
logger.info("LLM 配置已更新")
|
||
|
||
return {"success": True, "message": "LLM 配置已保存"}
|
||
|
||
|
||
@app.put("/api/config/channels")
|
||
async def update_channels_config(data: ChannelConfigUpdate) -> dict[str, Any]:
|
||
"""更新通讯渠道配置"""
|
||
settings = get_settings()
|
||
|
||
if data.wework:
|
||
if "enabled" in data.wework:
|
||
settings.channels.wework.enabled = data.wework["enabled"]
|
||
if data.wework.get("corp_id"):
|
||
settings.wework_corp_id = data.wework["corp_id"]
|
||
if data.wework.get("agent_id"):
|
||
settings.wework_agent_id = data.wework["agent_id"]
|
||
if data.wework.get("secret") and data.wework["secret"] != "***":
|
||
settings.wework_secret = data.wework["secret"]
|
||
if data.wework.get("token"):
|
||
settings.wework_token = data.wework["token"]
|
||
if data.wework.get("encoding_aes_key") and data.wework["encoding_aes_key"] != "***":
|
||
settings.wework_encoding_aes_key = data.wework["encoding_aes_key"]
|
||
|
||
if data.feishu:
|
||
if "enabled" in data.feishu:
|
||
settings.channels.feishu.enabled = data.feishu["enabled"]
|
||
if data.feishu.get("app_id"):
|
||
settings.feishu_app_id = data.feishu["app_id"]
|
||
if data.feishu.get("app_secret") and data.feishu["app_secret"] != "***":
|
||
settings.feishu_app_secret = data.feishu["app_secret"]
|
||
if data.feishu.get("verification_token"):
|
||
settings.feishu_verification_token = data.feishu["verification_token"]
|
||
if data.feishu.get("encrypt_key") and data.feishu["encrypt_key"] != "***":
|
||
settings.feishu_encrypt_key = data.feishu["encrypt_key"]
|
||
|
||
save_config(settings)
|
||
logger.info("通讯渠道配置已更新")
|
||
|
||
return {"success": True, "message": "通讯渠道配置已保存"}
|
||
|
||
|
||
@app.put("/api/config/proxy")
|
||
async def update_proxy_config(data: ProxyConfigUpdate) -> dict[str, Any]:
|
||
"""更新代理配置"""
|
||
settings = get_settings()
|
||
|
||
if data.enabled is not None:
|
||
settings.proxy.enabled = data.enabled
|
||
if data.http is not None:
|
||
settings.proxy.http = data.http
|
||
if data.https is not None:
|
||
settings.proxy.https = data.https
|
||
if data.no_proxy is not None:
|
||
settings.proxy.no_proxy = data.no_proxy
|
||
if data.auto_detect is not None:
|
||
settings.proxy.auto_detect = data.auto_detect
|
||
|
||
save_config(settings)
|
||
logger.info("代理配置已更新")
|
||
|
||
return {"success": True, "message": "代理配置已保存"}
|
||
|
||
|
||
@app.post("/api/llm/test/{provider}")
|
||
async def test_llm_connection(provider: str) -> dict[str, Any]:
|
||
"""测试 LLM 连接"""
|
||
try:
|
||
from minenasai.llm import get_llm_manager
|
||
|
||
manager = get_llm_manager()
|
||
# 发送一个简单的测试消息
|
||
response = await manager.chat(
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
provider=provider,
|
||
max_tokens=10
|
||
)
|
||
return {"success": True, "message": "连接成功"}
|
||
except Exception as e:
|
||
return {"success": False, "error": str(e)}
|
||
|
||
|
||
@app.post("/api/proxy/test")
|
||
async def test_proxy() -> dict[str, Any]:
|
||
"""测试代理连接"""
|
||
import httpx
|
||
|
||
settings = get_settings()
|
||
if not settings.proxy.enabled:
|
||
return {"success": False, "error": "代理未启用"}
|
||
|
||
try:
|
||
proxy_url = settings.proxy.http or settings.proxy.https
|
||
async with httpx.AsyncClient(proxy=proxy_url, timeout=10) as client:
|
||
response = await client.get("https://httpbin.org/ip")
|
||
return {"success": True, "ip": response.json().get("origin")}
|
||
except Exception as e:
|
||
return {"success": False, "error": str(e)}
|
||
|
||
|
||
@app.post("/api/proxy/detect")
|
||
async def detect_proxy() -> dict[str, Any]:
|
||
"""检测本地代理"""
|
||
import socket
|
||
|
||
common_ports = [7890, 7891, 1080, 1087, 10808]
|
||
|
||
for port in common_ports:
|
||
try:
|
||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
sock.settimeout(1)
|
||
result = sock.connect_ex(('127.0.0.1', port))
|
||
sock.close()
|
||
if result == 0:
|
||
proxy_url = f"http://127.0.0.1:{port}"
|
||
return {
|
||
"detected": True,
|
||
"http": proxy_url,
|
||
"https": proxy_url,
|
||
"port": port
|
||
}
|
||
except Exception:
|
||
continue
|
||
|
||
return {"detected": False}
|
||
|
||
|
||
# ==================== Agent 管理 API ====================
|
||
|
||
|
||
class AgentCreate(BaseModel):
|
||
"""Agent 创建/更新"""
|
||
id: str
|
||
name: str
|
||
workspace_path: str = "~/.config/minenasai/workspace"
|
||
model: str = "claude-sonnet-4-20250514"
|
||
temperature: float = 0.7
|
||
tools: dict[str, list[str]] | None = None
|
||
sandbox: dict[str, Any] | None = None
|
||
|
||
|
||
@app.get("/api/agents")
|
||
async def list_agents() -> dict[str, Any]:
|
||
"""列出所有 Agent"""
|
||
settings = get_settings()
|
||
agents = []
|
||
|
||
for agent in settings.agents.items:
|
||
agents.append({
|
||
"id": agent.id,
|
||
"name": agent.name,
|
||
"workspace_path": agent.workspace_path,
|
||
"model": settings.agents.default_model,
|
||
"temperature": settings.agents.temperature,
|
||
"tools": {
|
||
"allow": agent.tools.allow,
|
||
"deny": agent.tools.deny,
|
||
},
|
||
"sandbox": {
|
||
"mode": agent.sandbox.mode,
|
||
},
|
||
"status": "idle"
|
||
})
|
||
|
||
return {"agents": agents}
|
||
|
||
|
||
@app.post("/api/agents")
|
||
async def create_agent(data: AgentCreate) -> dict[str, Any]:
|
||
"""创建 Agent"""
|
||
from minenasai.core.config import AgentConfig, AgentToolsConfig, AgentSandboxConfig
|
||
|
||
settings = get_settings()
|
||
|
||
# 检查 ID 是否已存在
|
||
for agent in settings.agents.items:
|
||
if agent.id == data.id:
|
||
return JSONResponse(
|
||
status_code=400,
|
||
content={"success": False, "error": "Agent ID 已存在"}
|
||
)
|
||
|
||
# 创建新 Agent
|
||
tools_config = AgentToolsConfig(
|
||
allow=data.tools.get("allow", []) if data.tools else [],
|
||
deny=data.tools.get("deny", []) if data.tools else []
|
||
)
|
||
sandbox_config = AgentSandboxConfig(
|
||
mode=data.sandbox.get("mode", "workspace") if data.sandbox else "workspace"
|
||
)
|
||
|
||
new_agent = AgentConfig(
|
||
id=data.id,
|
||
name=data.name,
|
||
workspace_path=data.workspace_path,
|
||
tools=tools_config,
|
||
sandbox=sandbox_config
|
||
)
|
||
|
||
settings.agents.items.append(new_agent)
|
||
save_config(settings)
|
||
|
||
logger.info("Agent 已创建", agent_id=data.id)
|
||
return {"success": True, "message": "Agent 创建成功"}
|
||
|
||
|
||
@app.put("/api/agents/{agent_id}")
|
||
async def update_agent(agent_id: str, data: AgentCreate) -> dict[str, Any]:
|
||
"""更新 Agent"""
|
||
settings = get_settings()
|
||
|
||
for i, agent in enumerate(settings.agents.items):
|
||
if agent.id == agent_id:
|
||
agent.name = data.name
|
||
agent.workspace_path = data.workspace_path
|
||
if data.tools:
|
||
agent.tools.allow = data.tools.get("allow", [])
|
||
agent.tools.deny = data.tools.get("deny", [])
|
||
if data.sandbox:
|
||
agent.sandbox.mode = data.sandbox.get("mode", "workspace")
|
||
|
||
save_config(settings)
|
||
logger.info("Agent 已更新", agent_id=agent_id)
|
||
return {"success": True, "message": "Agent 更新成功"}
|
||
|
||
return JSONResponse(
|
||
status_code=404,
|
||
content={"success": False, "error": "Agent 不存在"}
|
||
)
|
||
|
||
|
||
@app.delete("/api/agents/{agent_id}")
|
||
async def delete_agent(agent_id: str) -> dict[str, Any]:
|
||
"""删除 Agent"""
|
||
settings = get_settings()
|
||
|
||
for i, agent in enumerate(settings.agents.items):
|
||
if agent.id == agent_id:
|
||
settings.agents.items.pop(i)
|
||
save_config(settings)
|
||
logger.info("Agent 已删除", agent_id=agent_id)
|
||
return {"success": True, "message": "Agent 删除成功"}
|
||
|
||
return JSONResponse(
|
||
status_code=404,
|
||
content={"success": False, "error": "Agent 不存在"}
|
||
)
|
||
|
||
|
||
# ==================== 定时任务 API ====================
|
||
|
||
|
||
class CronJobCreate(BaseModel):
|
||
"""定时任务创建/更新"""
|
||
name: str
|
||
agent_id: str = "main"
|
||
schedule: str
|
||
task: str
|
||
enabled: bool = True
|
||
|
||
|
||
@app.get("/api/cron-jobs")
|
||
async def list_cron_jobs() -> dict[str, Any]:
|
||
"""列出所有定时任务"""
|
||
# TODO: 从数据库获取
|
||
return {"jobs": []}
|
||
|
||
|
||
@app.post("/api/cron-jobs")
|
||
async def create_cron_job(data: CronJobCreate) -> dict[str, Any]:
|
||
"""创建定时任务"""
|
||
# TODO: 保存到数据库
|
||
logger.info("定时任务已创建", name=data.name)
|
||
return {"success": True, "message": "定时任务创建成功"}
|
||
|
||
|
||
@app.put("/api/cron-jobs/{job_id}")
|
||
async def update_cron_job(job_id: str, data: CronJobCreate) -> dict[str, Any]:
|
||
"""更新定时任务"""
|
||
# TODO: 更新数据库
|
||
return {"success": True, "message": "定时任务更新成功"}
|
||
|
||
|
||
@app.delete("/api/cron-jobs/{job_id}")
|
||
async def delete_cron_job(job_id: str) -> dict[str, Any]:
|
||
"""删除定时任务"""
|
||
# TODO: 从数据库删除
|
||
return {"success": True, "message": "定时任务删除成功"}
|
||
|
||
|
||
@app.post("/api/cron-jobs/{job_id}/toggle")
|
||
async def toggle_cron_job(job_id: str) -> dict[str, Any]:
|
||
"""切换定时任务启用状态"""
|
||
# TODO: 更新数据库
|
||
return {"success": True, "message": "状态已切换"}
|
||
|
||
|
||
@app.post("/api/cron-jobs/{job_id}/run")
|
||
async def run_cron_job(job_id: str) -> dict[str, Any]:
|
||
"""立即执行定时任务"""
|
||
# TODO: 触发任务执行
|
||
return {"success": True, "message": "任务已触发执行"}
|
||
|
||
|
||
# ==================== 日志 API ====================
|
||
|
||
|
||
@app.get("/api/logs")
|
||
async def get_logs(limit: int = 100) -> dict[str, Any]:
|
||
"""获取最近日志"""
|
||
# TODO: 从日志文件读取
|
||
return {"logs": []}
|
||
|
||
|
||
@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,
|
||
})
|