292 lines
9.4 KiB
Python
292 lines
9.4 KiB
Python
"""代理 & 白名单配置路由 + WebSocket 实时日志"""
|
||
|
||
import asyncio
|
||
import threading
|
||
import uuid
|
||
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
|
||
from sqlalchemy.orm import Session
|
||
|
||
from ..database import get_db
|
||
from ..models import User, ProxyConfig as ProxyConfigModel, AuditLog
|
||
from ..schemas import ProxyConfigOut, ProxyConfigUpdate
|
||
from ..deps import require_permission
|
||
from core.douyu.proxy import resolve_working_proxy, verify_proxy_url, parse_proxy_response
|
||
|
||
router = APIRouter(prefix="/api/proxy", tags=["代理与白名单"])
|
||
|
||
# 运行中的测试: test_id -> {log_queue, loop, result}
|
||
_active_tests: dict[str, dict] = {}
|
||
|
||
|
||
def _get_or_create(db: Session) -> ProxyConfigModel:
|
||
cfg = db.query(ProxyConfigModel).first()
|
||
if not cfg:
|
||
cfg = ProxyConfigModel()
|
||
db.add(cfg)
|
||
db.commit()
|
||
db.refresh(cfg)
|
||
return cfg
|
||
|
||
|
||
@router.get("", response_model=ProxyConfigOut)
|
||
def get_proxy_config(
|
||
db: Session = Depends(get_db),
|
||
_: User = Depends(require_permission("proxy:manage")),
|
||
):
|
||
return _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")),
|
||
):
|
||
cfg = _get_or_create(db)
|
||
cfg.enabled = req.enabled
|
||
cfg.api_url = req.api_url
|
||
cfg.http = req.http
|
||
cfg.https = req.https
|
||
cfg.whitelist_enabled = req.whitelist_enabled
|
||
cfg.whitelist_uid = req.whitelist_uid
|
||
cfg.whitelist_ukey = req.whitelist_ukey
|
||
db.commit()
|
||
db.refresh(cfg)
|
||
|
||
db.add(AuditLog(user_id=current.id, username=current.username,
|
||
action="proxy:update", target="proxy_config"))
|
||
db.commit()
|
||
return cfg
|
||
|
||
|
||
# ---- WebSocket 日志推送 ----
|
||
|
||
@router.websocket("/ws/test/{test_id}")
|
||
async def ws_test_logs(websocket: WebSocket, test_id: str):
|
||
"""WebSocket 推送代理/白名单测试实时日志。"""
|
||
await websocket.accept()
|
||
|
||
test = _active_tests.get(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:
|
||
_active_tests.pop(test_id, None)
|
||
|
||
|
||
# ---- 异步测试执行 ----
|
||
|
||
def _run_proxy_test(
|
||
cfg: ProxyConfigModel,
|
||
log_queue: asyncio.Queue,
|
||
loop: asyncio.AbstractEventLoop,
|
||
):
|
||
"""在线程中执行代理测试。"""
|
||
|
||
def push(level, message):
|
||
asyncio.run_coroutine_threadsafe(
|
||
log_queue.put({"level": level, "message": message}),
|
||
loop,
|
||
)
|
||
|
||
try:
|
||
if not cfg.enabled:
|
||
push("error", "代理未启用")
|
||
push("result", "")
|
||
return
|
||
|
||
# 静态代理
|
||
if cfg.http or cfg.https:
|
||
proxy_url = cfg.http or cfg.https
|
||
push("info", f"验证静态代理: {proxy_url}")
|
||
ok, msg = verify_proxy_url(proxy_url)
|
||
push("success" if ok else "error", msg)
|
||
push("result", "")
|
||
return
|
||
|
||
# API代理
|
||
if cfg.api_url:
|
||
whitelist_uid = cfg.whitelist_uid if cfg.whitelist_enabled else ""
|
||
whitelist_ukey = cfg.whitelist_ukey if cfg.whitelist_enabled else ""
|
||
|
||
proxy_url, msg = resolve_working_proxy(
|
||
api_url=cfg.api_url,
|
||
whitelist_uid=whitelist_uid,
|
||
whitelist_ukey=whitelist_ukey,
|
||
max_attempts=3,
|
||
log_func=push,
|
||
)
|
||
if proxy_url:
|
||
push("success", f"代理可用: {proxy_url}")
|
||
else:
|
||
push("error", msg)
|
||
push("result", "")
|
||
return
|
||
|
||
push("error", "未配置代理地址或API")
|
||
push("result", "")
|
||
except Exception as e:
|
||
push("error", f"测试异常: {e}")
|
||
push("result", "")
|
||
|
||
|
||
def _run_whitelist_test(
|
||
cfg: ProxyConfigModel,
|
||
log_queue: asyncio.Queue,
|
||
loop: asyncio.AbstractEventLoop,
|
||
):
|
||
"""在线程中执行白名单测试。"""
|
||
import requests as req_lib
|
||
import re
|
||
from core.douyu.whitelist import WhitelistManager, get_exit_ip_via_proxy
|
||
|
||
def push(level, message):
|
||
asyncio.run_coroutine_threadsafe(
|
||
log_queue.put({"level": level, "message": message}),
|
||
loop,
|
||
)
|
||
|
||
try:
|
||
if not cfg.whitelist_enabled:
|
||
push("error", "白名单未启用")
|
||
push("result", "")
|
||
return
|
||
if not cfg.whitelist_uid or not cfg.whitelist_ukey:
|
||
push("error", "未配置白名单UID/UKEY")
|
||
push("result", "")
|
||
return
|
||
|
||
manager = WhitelistManager(cfg.whitelist_uid, cfg.whitelist_ukey)
|
||
|
||
# 1. 测试API连接
|
||
push("info", "测试白名单API连接...")
|
||
ok, msg = manager.test_connection()
|
||
push("info" if ok else "error", f"白名单API: {msg}")
|
||
if not ok:
|
||
push("result", "")
|
||
return
|
||
|
||
# 2. 获取本机公网IP（白名单需要的是本机IP，不是代理出口IP）
|
||
push("info", "正在获取本机公网IP...")
|
||
|
||
# 优先：从代理API响应中提取（代理API返回"请先添加白名单：xxx"时，xxx就是本机IP）
|
||
local_ip = None
|
||
if cfg.api_url:
|
||
try:
|
||
resp = req_lib.get(cfg.api_url, timeout=10)
|
||
text = resp.text.strip()
|
||
push("info", f"代理API响应: {text[:80]}")
|
||
_, whitelist_ip = parse_proxy_response(text)
|
||
if whitelist_ip:
|
||
local_ip = whitelist_ip
|
||
push("info", f"从代理API获取到本机IP: {local_ip}")
|
||
except Exception as e:
|
||
push("warning", f"代理API请求失败: {e}")
|
||
|
||
# 备用：直接访问IP检测服务获取本机公网IP
|
||
if not local_ip:
|
||
push("info", "通过IP检测服务获取本机公网IP...")
|
||
for url in [
|
||
'https://qifu-api.baidubce.com/ip/local/geo/v1/district',
|
||
'https://myip.ipip.net',
|
||
'https://4.ipw.cn',
|
||
]:
|
||
try:
|
||
resp = req_lib.get(url, timeout=6, headers={"User-Agent": "Mozilla/5.0"})
|
||
match = re.search(r'(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})', resp.text)
|
||
if match:
|
||
local_ip = match.group(1)
|
||
break
|
||
except Exception:
|
||
continue
|
||
|
||
if not local_ip:
|
||
push("error", "无法获取本机公网IP")
|
||
push("result", "")
|
||
return
|
||
|
||
push("info", f"本机公网IP: {local_ip}")
|
||
|
||
# 3. 检查并同步白名单
|
||
records = manager.get_whitelist_json()
|
||
in_list = any(r.get('IP') == local_ip for r in records)
|
||
push("info", f"白名单共 {len(records)} 条记录")
|
||
|
||
if in_list:
|
||
record = next((r for r in records if r.get('IP') == local_ip), {})
|
||
memo = record.get('MEMO', '')
|
||
if memo == manager.memo:
|
||
push("success", f"本机IP {local_ip} 已在白名单中 (备注正确)")
|
||
else:
|
||
push("warning", f'本机IP {local_ip} 备注不匹配 (当前: "{memo}")，更新中...')
|
||
sync_ok, sync_msg = manager.sync_ip(local_ip)
|
||
push("success" if sync_ok else "error", f"白名单更新: {sync_msg}")
|
||
else:
|
||
push("info", f"正在将 {local_ip} 添加到白名单...")
|
||
sync_ok, sync_msg = manager.sync_ip(local_ip)
|
||
push("success" if sync_ok else "error", f"白名单同步: {sync_msg}")
|
||
|
||
push("result", "")
|
||
except Exception as e:
|
||
push("error", f"测试异常: {e}")
|
||
push("result", "")
|
||
|
||
|
||
# ---- API 端点 ----
|
||
|
||
@router.post("/test")
|
||
async def test_proxy(
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(require_permission("proxy:manage")),
|
||
):
|
||
"""启动代理测试（异步执行，通过 WebSocket 推送日志）。"""
|
||
cfg = _get_or_create(db)
|
||
|
||
test_id = uuid.uuid4().hex[:12]
|
||
log_queue = asyncio.Queue()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
_active_tests[test_id] = {"log_queue": log_queue, "loop": loop}
|
||
|
||
thread = threading.Thread(target=_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 = _get_or_create(db)
|
||
|
||
test_id = uuid.uuid4().hex[:12]
|
||
log_queue = asyncio.Queue()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
_active_tests[test_id] = {"log_queue": log_queue, "loop": loop}
|
||
|
||
thread = threading.Thread(target=_run_whitelist_test, args=(cfg, log_queue, loop), daemon=True)
|
||
thread.start()
|
||
|
||
return {"test_id": test_id, "success": True}
|