138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
"""登录任务路由 + WebSocket 实时日志"""
|
|
|
|
import asyncio
|
|
import threading
|
|
from datetime import datetime
|
|
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..database import get_db, SessionLocal
|
|
from ..models import User, Account, LoginTask, ProxyConfig as ProxyConfigModel
|
|
from ..schemas import LoginBatchRequest, LoginTaskOut
|
|
from ..deps import get_current_user, require_permission
|
|
from ..permissions import has_permission
|
|
from ..services.login_service import LoginBatchRunner
|
|
|
|
router = APIRouter(prefix="/api/login", tags=["登录任务"])
|
|
|
|
# 运行中的批次: batch_id -> {runner, log_queue, loop}
|
|
_active_batches: dict[str, dict] = {}
|
|
|
|
|
|
@router.post("/batch")
|
|
def create_batch(
|
|
req: LoginBatchRequest,
|
|
db: Session = Depends(get_db),
|
|
current: User = Depends(require_permission("login:batch")),
|
|
):
|
|
"""创建批量登录任务。"""
|
|
if not req.account_ids:
|
|
raise HTTPException(status_code=400, detail="请选择账号")
|
|
|
|
# 读取代理配置
|
|
proxy = db.query(ProxyConfigModel).first()
|
|
|
|
# 权限过滤账号
|
|
valid_ids = []
|
|
for aid in req.account_ids:
|
|
acc = db.query(Account).filter(Account.id == aid).first()
|
|
if not acc:
|
|
continue
|
|
if not has_permission(current.role, "login:view_all"):
|
|
if acc.assigned_to != current.id:
|
|
continue
|
|
valid_ids.append(aid)
|
|
|
|
if not valid_ids:
|
|
raise HTTPException(status_code=403, detail="没有可登录的账号")
|
|
|
|
# 创建执行器（用独立的 DB 会话，因为在线程中运行）
|
|
thread_db = SessionLocal()
|
|
runner = LoginBatchRunner(
|
|
db=thread_db,
|
|
account_ids=valid_ids,
|
|
created_by=current.id,
|
|
creator_role=current.role,
|
|
max_geetest_retries=req.max_geetest_retries,
|
|
proxy_config=proxy,
|
|
)
|
|
|
|
batch_id = runner.batch_id
|
|
|
|
# 启动线程
|
|
thread = threading.Thread(target=runner.run, daemon=True)
|
|
thread.start()
|
|
|
|
return {"batch_id": batch_id, "count": len(valid_ids), "success": True}
|
|
|
|
|
|
@router.get("/tasks", response_model=list[LoginTaskOut])
|
|
def list_tasks(
|
|
batch_id: str | None = None,
|
|
db: Session = Depends(get_db),
|
|
current: User = Depends(get_current_user),
|
|
):
|
|
"""查看登录任务列表。"""
|
|
query = db.query(LoginTask)
|
|
|
|
# 客服只能看自己账号的任务
|
|
if not has_permission(current.role, "login:view_all"):
|
|
query = query.join(Account, LoginTask.account_id == Account.id).filter(
|
|
Account.assigned_to == current.id
|
|
)
|
|
|
|
if batch_id:
|
|
query = query.filter(LoginTask.batch_id == batch_id)
|
|
|
|
tasks = query.order_by(LoginTask.id.desc()).limit(200).all()
|
|
result = []
|
|
for t in tasks:
|
|
acc = db.query(Account).filter(Account.id == t.account_id).first()
|
|
result.append(LoginTaskOut(
|
|
id=t.id, batch_id=t.batch_id, account_id=t.account_id,
|
|
account_username=acc.username if acc else "",
|
|
status=t.status, cookie=t.cookie or "", message=t.message or "",
|
|
created_by=t.created_by, created_at=t.created_at, finished_at=t.finished_at,
|
|
))
|
|
return result
|
|
|
|
|
|
@router.post("/stop/{batch_id}")
|
|
def stop_batch(
|
|
batch_id: str,
|
|
current: User = Depends(require_permission("login:batch")),
|
|
):
|
|
batch = _active_batches.get(batch_id)
|
|
if batch:
|
|
batch["runner"].stop()
|
|
return {"message": "已发送停止信号", "success": True}
|
|
raise HTTPException(status_code=404, detail="批次不存在或已结束")
|
|
|
|
|
|
@router.websocket("/ws/login/{batch_id}")
|
|
async def ws_login_logs(websocket: WebSocket, batch_id: str):
|
|
"""WebSocket 推送登录实时日志。"""
|
|
await websocket.accept()
|
|
|
|
log_queue = asyncio.Queue()
|
|
loop = asyncio.get_event_loop()
|
|
|
|
# 查找已运行的批次，或等待新批次
|
|
# 简化：直接把 log_queue 注册到全局，前端创建批次后连 ws
|
|
_active_batches[batch_id] = {
|
|
"log_queue": log_queue,
|
|
"loop": loop,
|
|
}
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
msg = await asyncio.wait_for(log_queue.get(), timeout=30)
|
|
await websocket.send_json(msg)
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json({"level": "heartbeat", "message": ""})
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
_active_batches.pop(batch_id, None)
|