147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
"""代理 & 白名单配置路由 + WebSocket 实时日志"""
|
|
|
|
import asyncio
|
|
import threading
|
|
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..database import get_db
|
|
from ..models import User
|
|
from ..schemas import ProxyConfigOut, ProxyConfigUpdate, PlatformInfo, PlatformFieldDef
|
|
from ..deps import require_permission, authenticate_websocket
|
|
from ..services.proxy_service import proxy_service
|
|
from core.douyu.proxy_platforms import (
|
|
get_platform_names,
|
|
get_platform_labels,
|
|
get_credential_fields,
|
|
)
|
|
|
|
router = APIRouter(prefix="/api/proxy", tags=["代理与白名单"])
|
|
|
|
|
|
@router.get("", response_model=ProxyConfigOut)
|
|
def get_proxy_config(
|
|
db: Session = Depends(get_db),
|
|
_: User = Depends(require_permission("proxy:manage")),
|
|
):
|
|
return proxy_service.get_or_create(db)
|
|
|
|
|
|
@router.put("", response_model=ProxyConfigOut)
|
|
def update_proxy_config(
|
|
req: ProxyConfigUpdate,
|
|
db: Session = Depends(get_db),
|
|
current: User = Depends(require_permission("proxy:manage")),
|
|
):
|
|
return proxy_service.update_config(
|
|
db,
|
|
enabled=req.enabled if req.enabled is not None else False,
|
|
api_url=req.api_url if req.api_url is not None else "",
|
|
http=req.http if req.http is not None else "",
|
|
https=req.https if req.https is not None else "",
|
|
whitelist_enabled=req.whitelist_enabled
|
|
if req.whitelist_enabled is not None
|
|
else False,
|
|
whitelist_platform=req.whitelist_platform or "xiequ",
|
|
whitelist_credentials=req.whitelist_credentials,
|
|
whitelist_uid=req.whitelist_uid,
|
|
whitelist_ukey=req.whitelist_ukey,
|
|
current_user=current,
|
|
)
|
|
|
|
|
|
@router.get("/platforms", response_model=list[PlatformInfo])
|
|
def list_platforms():
|
|
"""获取所有可用的代理白名单平台及其凭据字段定义。"""
|
|
labels = get_platform_labels()
|
|
result = []
|
|
for name in get_platform_names():
|
|
fields = get_credential_fields(name)
|
|
result.append(
|
|
PlatformInfo(
|
|
name=name,
|
|
label=labels.get(name, name),
|
|
credential_fields=[PlatformFieldDef(**f) for f in fields],
|
|
)
|
|
)
|
|
return result
|
|
|
|
|
|
# ---- WebSocket 日志推送 ----
|
|
|
|
|
|
@router.websocket("/ws/test/{test_id}")
|
|
async def ws_test_logs(websocket: WebSocket, test_id: str):
|
|
"""WebSocket 推送代理/白名单测试实时日志（需认证）。"""
|
|
user = authenticate_websocket(websocket)
|
|
if not user:
|
|
await websocket.close(code=1008, reason="未授权")
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
test = proxy_service.get_test(test_id)
|
|
if not test:
|
|
await websocket.send_json({"level": "error", "message": "测试任务不存在"})
|
|
await websocket.close()
|
|
return
|
|
|
|
log_queue: asyncio.Queue = test["log_queue"]
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
msg = await asyncio.wait_for(log_queue.get(), timeout=30)
|
|
await websocket.send_json(msg)
|
|
# 收到 result 消息表示测试结束
|
|
if msg.get("level") == "result":
|
|
await asyncio.sleep(0.1)
|
|
break
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json({"level": "heartbeat", "message": ""})
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
proxy_service.remove_test(test_id)
|
|
|
|
|
|
# ---- 异步测试 API 端点 ----
|
|
|
|
|
|
@router.post("/test")
|
|
async def test_proxy(
|
|
db: Session = Depends(get_db),
|
|
current: User = Depends(require_permission("proxy:manage")),
|
|
):
|
|
"""启动代理测试（异步执行，通过 WebSocket 推送日志）。"""
|
|
cfg = proxy_service.get_or_create(db)
|
|
test_id, log_queue, loop = proxy_service.register_test()
|
|
|
|
thread = threading.Thread(
|
|
target=proxy_service.run_proxy_test,
|
|
args=(cfg, log_queue, loop),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
|
|
return {"test_id": test_id, "success": True}
|
|
|
|
|
|
@router.post("/whitelist/test")
|
|
async def test_whitelist(
|
|
db: Session = Depends(get_db),
|
|
current: User = Depends(require_permission("whitelist:test")),
|
|
):
|
|
"""启动白名单测试（异步执行，通过 WebSocket 推送日志）。"""
|
|
cfg = proxy_service.get_or_create(db)
|
|
test_id, log_queue, loop = proxy_service.register_test()
|
|
|
|
thread = threading.Thread(
|
|
target=proxy_service.run_whitelist_test,
|
|
args=(cfg, log_queue, loop),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
|
|
return {"test_id": test_id, "success": True}
|