"""登录任务路由 + WebSocket 实时日志""" import asyncio import threading from fastapi import ( APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, ) from sqlalchemy import and_, func, or_ from sqlalchemy.orm import Session, load_only from ..database import SessionLocal, get_db from ..deps import authenticate_websocket, get_current_user, require_permission from ..models import Account, LoginTask, User from ..models import ProxyConfig as ProxyConfigModel from ..permissions import get_user_permissions, user_has_permission from ..schemas import LoginBatchRequest, LoginTaskOut from ..services.audit_service import record_audit from ..services.login_service import LoginBatchRunner, batch_registry router = APIRouter(prefix="/api/login", tags=["登录任务"]) def _effective_attempt_status(): """页面展示用的状态：优先最近一次登录尝试，没有尝试时回退到 CK 状态。 ``attempt_*`` 才是「这次登录怎么样了」，``status`` 只表示这条记录是否 持有可用 CK —— 两者分开，停止/失败/重启都不会让账号从列表里消失。 """ return func.coalesce( func.nullif(LoginTask.attempt_status, ""), LoginTask.status, ) def _login_task_summary(query): """按状态汇总登录任务，避免概览页拉完整任务列表。""" rows = ( query.order_by(None) .with_entities(_effective_attempt_status(), func.count(LoginTask.id)) .group_by(_effective_attempt_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") and 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() # 审计：批次是谁在什么时候发起、多少个账号、什么模式。 # 自动重登与人工登录共用同一条链路，没有这条记录就无法追溯来源。 try: record_audit( db, current, action="login:batch", target=batch_id, detail={ "mode": req.mode, "requested": len(req.account_ids), "accepted": len(valid_ids), "concurrency": req.concurrency, "api_strategy": req.api_strategy, }, ) db.commit() except Exception: # noqa: BLE001 审计失败不应影响已启动的批次 db.rollback() return {"batch_id": batch_id, "count": len(valid_ids), "success": True} def _visible_tasks_query(db: Session, current: User, batch_id: str | None): """按权限与批次过滤登录任务。客服只能看自己名下账号的任务。 这里已经**不需要 defer cookie** —— 登录记录不再持有凭证（2026-10-04 CK 解耦）。 """ 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: # 记录持有的 CK 属于产出它的批次，最近一次尝试属于新批次，两者都要算。 query = query.filter( or_(LoginTask.batch_id == batch_id, LoginTask.attempt_batch_id == batch_id) ) return query def _serialize_tasks(db: Session, rows) -> list[LoginTaskOut]: """把任务行拼上账号名后序列化。""" account_ids = [t.account_id for t in rows] accounts_map = {} if account_ids: accs = ( db.query(Account) .options(load_only(Account.id, Account.username)) .filter(Account.id.in_(account_ids)) .all() ) accounts_map = {a.id: a.username for a in accs} result = [] for t in rows: attempt_status = t.attempt_status or "" result.append( LoginTaskOut( id=t.id, batch_id=t.attempt_batch_id or t.batch_id, account_id=t.account_id, account_username=accounts_map.get(t.account_id, ""), status=attempt_status or t.status, cookie="", message=(t.attempt_message if attempt_status else "") or t.message or "", created_by=t.created_by, created_at=t.created_at, finished_at=t.attempt_at or t.finished_at, ) ) return result @router.get("/tasks") def list_tasks( batch_id: str | None = None, # 默认值保持 None（而不是 Query(None)）：这个接口会被测试与内部代码直接 # 调用，用 Query() 做默认值时拿到的会是那个 Query 对象而非 None。 page: int | None = None, page_size: int = 50, status: str | None = None, db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """查看登录任务列表。 **不传 `page` 时返回裸数组**（上限 200），沿用老调用方（Cookie 页等）的 契约；传了 `page` 才返回 ``{items, total, page, page_size}``，供表格分页。 传 `page` 时同时返回 `total`，是为了让前端能「全选所有筛选结果」—— 否则跨页选择无从下手，而任务量早就过了单页能装下的规模。 """ query = _visible_tasks_query(db, current, batch_id) if status: # 前端看到的是「最近一次尝试」的状态，与列表序列化口径保持一致。 query = query.filter( or_( and_( LoginTask.attempt_status != "", LoginTask.attempt_status == status ), and_(LoginTask.attempt_status == "", LoginTask.status == status), ) ) ordered = query.order_by(LoginTask.id.desc()) if page is None: return _serialize_tasks(db, ordered.limit(200).all()) page = max(1, int(page)) page_size = max(1, min(int(page_size), 500)) total = ordered.count() rows = ordered.offset((page - 1) * page_size).limit(page_size).all() return { "items": _serialize_tasks(db, rows), "total": total, "page": page, "page_size": page_size, } @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) def _in_flight_clause(): """不可删除的记录：正被批次处理的**在途尝试**。 ⚠️ **CK 自 2026-10-04 起不再是护栏理由** —— CK 已迁到 `accounts.cookie`， 删登录历史不再影响任何账号的登录态（这正是解耦的目的）。但「在途」必须保留： 删掉会让批次线程找不到记录；停批次请走 ``/stop/{batch_id}``。 """ in_flight = ("pending", "running") return or_( and_( LoginTask.attempt_status != "", LoginTask.attempt_status.in_(in_flight), ), and_(LoginTask.attempt_status == "", LoginTask.status.in_(in_flight)), ) def _deletable_clause(): """可安全删除的记录：不在途。""" return ~_in_flight_clause() def _is_in_flight(task: LoginTask) -> bool: """单条记录是否在途（与 `_in_flight_clause` 同口径）。""" in_flight = ("pending", "running") if task.attempt_status: return task.attempt_status in in_flight return (task.status or "") in in_flight @router.delete("/tasks/{task_id}") def delete_task( task_id: int, db: Session = Depends(get_db), current: User = Depends(require_permission("login:batch")), ): """删除单个登录任务（**在途的记录不允许删除**；删历史不影响 CK）。""" task = db.query(LoginTask).filter(LoginTask.id == task_id).first() if not task: raise HTTPException(status_code=404, detail="任务不存在") if _is_in_flight(task): raise HTTPException( status_code=409, detail="该记录正在被批次处理，删除会让批次失去目标；请先停止批次", ) db.delete(task) db.commit() return {"message": "已删除", "success": True} @router.delete("/tasks") def delete_tasks( task_ids: str = "", batch_id: str | None = None, status: str | None = None, all_matching: bool = False, db: Session = Depends(get_db), current: User = Depends(require_permission("login:batch")), ): """批量删除登录任务。三种模式，互斥： 1. ``task_ids``：删指定 id（表格勾选）； 2. ``all_matching=true``：删当前筛选条件下的**全部**（跨页全选用）； 3. 两者都没给 → 400，绝不默认全删。 两条护栏（都会在响应里报告跳过数量）： - **在途（pending/running）的记录不删**：删掉会让批次线程找不到记录； 停批次请走 ``/stop/{batch_id}``。 - ~~持有 CK 的记录不删~~ —— **这条护栏已随 CK 解耦（2026-10-04）取消**： CK 现在挂在账号上，删登录历史不影响登录态。历史终于可以自由清理。 """ # 护栏：不删在途任务，避免批次线程与数据库状态打架。 not_in_flight = _deletable_clause() if task_ids: ids = [int(x) for x in task_ids.split(",") if x.strip().isdigit()] if not ids: raise HTTPException(status_code=400, detail="无效的任务ID") elif all_matching: # 带 join 的 query 不能直接 bulk delete（SQLAlchemy 会拒绝）， # 所以先按权限与筛选条件查出 id，再按主键删。 query = _visible_tasks_query(db, current, batch_id) if status: query = query.filter( or_( and_( LoginTask.attempt_status != "", LoginTask.attempt_status == status, ), and_(LoginTask.attempt_status == "", LoginTask.status == status), ) ) ids = [row[0] for row in query.with_entities(LoginTask.id).all()] else: raise HTTPException(status_code=400, detail="请指定任务ID或 all_matching=true") if not ids: return { "message": "没有符合条件的任务", "deleted": 0, "skipped": 0, "success": True, } rows = ( db.query(LoginTask.id, not_in_flight.label("ok_to_delete")) .filter(LoginTask.id.in_(ids)) .all() ) target_ids = [r[0] for r in rows if r[1]] skipped = len(rows) - len(target_ids) if not target_ids: return { "message": f"没有可删除的记录（{skipped} 条在途，已跳过）", "deleted": 0, "skipped": skipped, "success": True, } deleted = ( db.query(LoginTask) .filter(LoginTask.id.in_(target_ids)) .delete(synchronize_session=False) ) db.commit() message = f"已删除 {deleted} 个任务" if skipped: message += f"（{skipped} 条在途，已跳过）" return {"message": message, "deleted": deleted, "skipped": skipped, "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 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)