"""登录任务路由 + WebSocket 实时日志""" import asyncio import threading from datetime import datetime from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect from sqlalchemy import func from sqlalchemy.orm import Session, defer 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=["登录任务"]) def _login_task_summary(query): """按状态汇总登录任务，避免概览页拉完整任务列表。""" rows = ( query.order_by(None) .with_entities(LoginTask.status, func.count(LoginTask.id)) .group_by(LoginTask.status) .all() ) status_counts = {status or "": count for status, count in rows} return { "total": sum(status_counts.values()), "success": status_counts.get("success", 0), "failed": sum(status_counts.get(status, 0) for status in ("failed", "error")), "status_counts": status_counts, } @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 = [] seen_ids = set() for aid in req.account_ids: if aid in seen_ids: continue seen_ids.add(aid) 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: action_name = "检测" if req.mode == "check" else "登录" raise HTTPException(status_code=403, detail=f"没有可{action_name}的账号") # 在主事件循环中创建 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_login_retries=req.max_login_retries, max_total_time=req.max_total_time, proxy_config=proxy, log_queue=log_queue, loop=loop, concurrency=req.concurrency, api_strategy=req.api_strategy, mode=req.mode, ) 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).options(defer(LoginTask.cookie)).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="", message=t.message or "", created_by=t.created_by, created_at=t.created_at, finished_at=t.finished_at, )) return result @router.get("/tasks/summary") def tasks_summary( 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) return _login_task_summary(query) @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: if batch.get("finished"): # 已结束但 WS 尚未清理（或从未有 WS 订阅）：回收注册表条目。 batch_registry.pop(batch_id) raise HTTPException(status_code=404, detail="批次已结束") 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 # 归属校验：非 view_all 只能订阅自己创建的批次（与虎牙批次一致）。 if not user_has_permission(user, "login:view_all"): check_db = SessionLocal() try: owned = ( check_db.query(LoginTask.id) .filter(LoginTask.batch_id == batch_id, LoginTask.created_by == user.id) .first() ) finally: check_db.close() if not owned: 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: # 只在批次真正结束后清理注册表，客户端断线/刷新页面后重连仍可继续订阅。 latest = batch_registry.get(batch_id) if latest and latest.get("finished"): batch_registry.pop(batch_id)