"""登录任务路由 + 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, authenticate_websocket from ..permissions import user_has_permission, get_user_permissions from ..services.login_service import LoginBatchRunner, batch_registry router = APIRouter(prefix="/api/login", tags=["登录任务"]) @router.post("/batch") async 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 user_has_permission(current, "login:view_all"): if acc.assigned_to != current.id: continue valid_ids.append(aid) if not valid_ids: raise HTTPException(status_code=403, detail="没有可登录的账号") # 在主事件循环中创建 log_queue,传给后台线程 log_queue = asyncio.Queue() loop = asyncio.get_running_loop() # 创建执行器(用独立的 DB 会话,因为在线程中运行) thread_db = SessionLocal() runner = LoginBatchRunner( db=thread_db, account_ids=valid_ids, created_by=current.id, creator_permissions=get_user_permissions(current), max_geetest_retries=req.max_geetest_retries, max_proxy_retries=req.max_proxy_retries, proxy_config=proxy, log_queue=log_queue, loop=loop, concurrency=req.concurrency, ) batch_id = runner.batch_id # 先注册到批次注册表,再启动线程,确保 WebSocket 连接时能找到 batch_registry.register(batch_id, log_queue, loop, runner) # 启动线程 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).join(Account, LoginTask.account_id == Account.id) # 客服只能看自己账号的任务 if not user_has_permission(current, "login:view_all"): query = query.filter(Account.assigned_to == current.id) if batch_id: query = query.filter(LoginTask.batch_id == batch_id) rows = query.order_by(LoginTask.id.desc()).limit(200).all() # 批量收集 account_id,一次性查出 username account_ids = [t.account_id for t in rows] accounts_map = {} if account_ids: accs = db.query(Account).filter(Account.id.in_(account_ids)).all() accounts_map = {a.id: a.username for a in accs} result = [] for t in rows: result.append(LoginTaskOut( id=t.id, batch_id=t.batch_id, account_id=t.account_id, account_username=accounts_map.get(t.account_id, ""), 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.delete("/tasks/{task_id}") def delete_task( task_id: int, db: Session = Depends(get_db), current: User = Depends(require_permission("login:batch")), ): """删除单个登录任务。""" task = db.query(LoginTask).filter(LoginTask.id == task_id).first() if not task: raise HTTPException(status_code=404, detail="任务不存在") db.delete(task) db.commit() return {"message": "已删除", "success": True} @router.delete("/tasks") def delete_tasks( task_ids: str = "", db: Session = Depends(get_db), current: User = Depends(require_permission("login:batch")), ): """批量删除登录任务。""" if not task_ids: raise HTTPException(status_code=400, detail="请指定任务ID") ids = [int(x) for x in task_ids.split(",") if x.strip().isdigit()] if not ids: raise HTTPException(status_code=400, detail="无效的任务ID") deleted = db.query(LoginTask).filter(LoginTask.id.in_(ids)).delete(synchronize_session=False) db.commit() return {"message": f"已删除 {deleted} 个任务", "deleted": deleted, "success": True} @router.post("/stop/{batch_id}") def stop_batch( batch_id: str, current: User = Depends(require_permission("login:batch")), ): batch = batch_registry.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 推送登录实时日志(需认证)。""" # 认证:从 cookie 或 token query param 验证用户身份 user = authenticate_websocket(websocket) if not user: await websocket.close(code=1008, reason="未授权") return await websocket.accept() # 从已注册的批次中获取 log_queue(由 create_batch 创建) batch = batch_registry.get(batch_id) if not batch: await websocket.send_json({"level": "error", "message": "批次不存在或已结束"}) await websocket.close() return log_queue: asyncio.Queue = batch["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: batch_registry.pop(batch_id)