"""Cookie 管理路由""" import csv import io import threading from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import UTC, datetime from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse from sqlalchemy import case, func, or_ from sqlalchemy.orm import Session, defer, joinedload from ..database import SessionLocal, get_db from ..deps import get_current_user, require_permission from ..models import Account, AuditLog, LoginTask, User from ..models import ProxyConfig as ProxyConfigModel from ..permissions import get_user_permissions, user_has_permission from ..schemas import CookieReloginRequest from ..services.cookie_check_service import check_douyu_cookie from ..services.login_service import BatchRegistry, LoginBatchRunner, get_relogin_limits def _fmt_dt(dt) -> str | None: """将 datetime 格式化为带时区的 ISO 字符串。""" if dt is None: return None if dt.tzinfo is None: dt = dt.replace(tzinfo=UTC) return dt.isoformat() router = APIRouter(prefix="/api/cookies", tags=["Cookie管理"]) cookie_relogin_registry = BatchRegistry() def _require_cookie_operation_perm(current: User) -> None: """允许脱敏的 CK 检测与重登,不授予 CK 查看或导出能力。""" if not ( user_has_permission(current, "cookie:operate") or user_has_permission(current, "cookie:view") ): raise HTTPException(status_code=403, detail="无权限: cookie:operate") def _parse_account_names(raw_names: str) -> list[str]: """解析前端粘贴的 Excel 账号名,每行一个并去重。""" names = [] seen = set() for value in ( (raw_names or "").replace("\r\n", "\n").replace("\r", "\n").split("\n") ): name = value.strip() if name and name not in seen: names.append(name) seen.add(name) return names def _order_cookie_tasks(query, selected_names: list[str]): """按自定义账号名的输入顺序排序,未指定时保持默认的最新优先。""" if selected_names: input_order = case( {name: index for index, name in enumerate(selected_names)}, value=Account.username, else_=len(selected_names), ) return query.order_by( input_order, LoginTask.finished_at.desc(), LoginTask.id.desc() ) return query.order_by(LoginTask.finished_at.desc(), LoginTask.id.desc()) def _check_one_cookie(task: LoginTask) -> dict: """将共享检测结果补充为 Cookie 记录的接口返回格式。""" result = check_douyu_cookie(task.cookie or "") result["id"] = task.id result["checked_at"] = result["checked_at"].isoformat() return result def _visible_cookie_tasks_query(db: Session, current: User): """返回当前用户可见的成功 Cookie 任务查询。""" query = db.query(LoginTask).filter(LoginTask.status == "success") # 无全量登录任务权限时,只能操作分配给自己的账号 Cookie。 if not user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id).filter( Account.assigned_to == current.id ) return query def _visible_cookie_operation_tasks_query(db: Session, current: User): """返回可检测/重登的 Cookie 记录,重登中或失败时仍保留在操作列表。""" query = db.query(LoginTask).filter( LoginTask.cookie != "", LoginTask.status.in_( ("success", "relogin_pending", "relogin_running", "relogin_failed") ), ) if not user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id).filter( Account.assigned_to == current.id ) return query @router.get("") def list_cookies( search: str = Query(""), tag: str = Query(""), account_names: str = Query(""), page: int | None = Query(None, ge=1), page_size: int = Query(20, ge=1, le=200), include_cookie: bool = Query(True), db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """查看登录成功的 Cookie 列表。""" if not user_has_permission(current, "cookie:view"): raise HTTPException(status_code=403, detail="无权限: cookie:view") if not include_cookie: query = _visible_cookie_tasks_query(db, current).options( defer(LoginTask.cookie) ) else: query = _visible_cookie_tasks_query(db, current).options( joinedload(LoginTask.account).joinedload(Account.assigned_user), ) search_text = (search or "").strip() tag_value = (tag or "").strip() selected_names = _parse_account_names(account_names) can_view_all = user_has_permission(current, "login:view_all") account_joined = not can_view_all if selected_names: if not account_joined: query = query.join(Account, LoginTask.account_id == Account.id) account_joined = True query = query.filter(Account.username.in_(selected_names)) if tag_value: if not account_joined: query = query.join(Account, LoginTask.account_id == Account.id) account_joined = True query = query.filter(Account.tag == tag_value) if search_text: pattern = f"%{search_text}%" if not account_joined: query = query.join(Account, LoginTask.account_id == Account.id) account_joined = True query = query.outerjoin(User, Account.assigned_to == User.id).filter( or_( Account.username.ilike(pattern), Account.tag.ilike(pattern), User.username.ilike(pattern), ) ) total = None if page is not None: total = ( query.order_by(None).with_entities(func.count(LoginTask.id)).scalar() or 0 ) query = _order_cookie_tasks(query, selected_names) if page is not None: query = query.offset((page - 1) * page_size).limit(page_size) tasks = query.all() # 批量查账号,避免 N+1 查询 account_ids = [t.account_id for t in tasks] accounts_map = {} if account_ids: account_query = db.query(Account).filter(Account.id.in_(account_ids)) if not include_cookie: account_query = account_query.options( joinedload(Account.assigned_user), defer(Account.password), defer(Account.email), defer(Account.email_password), ) accs = account_query.all() accounts_map = {a.id: a for a in accs} result = [] for t in tasks: acc = accounts_map.get(t.account_id) item = { "id": t.id, "batch_id": t.batch_id, "account_id": t.account_id, "account_username": acc.username if acc else "", "assigned_to": acc.assigned_to if acc else None, "assigned_username": acc.assigned_user.username if acc and acc.assigned_user else None, "created_at": _fmt_dt(t.finished_at), "ck_check_status": t.ck_check_status or "", "ck_check_result": t.ck_check_result, "ck_checked_at": _fmt_dt(t.ck_checked_at), } # 分页列表默认只返回预览,复制/导出时再获取完整敏感字段。 if include_cookie and user_has_permission(current, "cookie:view"): cookie = t.cookie or "" item["cookie"] = cookie item["cookie_preview"] = cookie[:50] + "..." if len(cookie) > 50 else cookie item["account_password"] = acc.password if acc else "" else: item["cookie"] = "" item["cookie_preview"] = "***" item["account_password"] = "" result.append(item) if page is not None: return { "items": result, "total": total or 0, "page": page, "page_size": page_size, } return result @router.get("/summary") def cookies_summary( db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """Cookie 管理统计,避免前端为了卡片统计拉全量 Cookie。""" if not user_has_permission(current, "cookie:view"): raise HTTPException(status_code=403, detail="无权限: cookie:view") query = _visible_cookie_tasks_query(db, current) if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) total = query.with_entities(func.count(LoginTask.id)).scalar() or 0 assigned_count = ( query.filter(Account.assigned_to.isnot(None)) .with_entities(func.count(LoginTask.id)) .scalar() or 0 ) return { "total": total, "assigned_count": assigned_count, "unassigned_count": max(0, total - assigned_count), } @router.get("/operations") def list_cookie_operations( search: str = Query(""), tag: str = Query(""), page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=200), db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """脱敏的 CK 操作列表,只用于检测与重登。""" _require_cookie_operation_perm(current) query = _visible_cookie_operation_tasks_query(db, current) if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) search_text = (search or "").strip() tag_value = (tag or "").strip() if tag_value: query = query.filter(Account.tag == tag_value) if search_text: query = query.filter( or_( Account.username.ilike(f"%{search_text}%"), Account.tag.ilike(f"%{search_text}%"), ) ) total = query.order_by(None).with_entities(func.count(LoginTask.id)).scalar() or 0 tasks = ( query.order_by(LoginTask.finished_at.desc(), LoginTask.id.desc()) .offset((page - 1) * page_size) .limit(page_size) .all() ) return { "items": [ { "id": task.id, "account_username": task.account.username if task.account else "", "tag": task.account.tag if task.account else "", "ck_check_status": task.ck_check_status or "", "ck_checked_at": _fmt_dt(task.ck_checked_at), "created_at": _fmt_dt(task.finished_at), "relogin_status": task.status if task.status != "success" else "", "relogin_message": task.message or "", "relogin_batch_id": task.batch_id if task.status in {"relogin_pending", "relogin_running"} else "", } for task in tasks ], "total": total, "page": page, "page_size": page_size, } @router.get("/operations/tags") def list_cookie_operation_tags( db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """返回当前用户可操作 CK 记录对应的标签,不跨越账号分配范围。""" _require_cookie_operation_perm(current) query = _visible_cookie_operation_tasks_query(db, current) if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) rows = ( query.filter(Account.tag != "", Account.tag.isnot(None)) .with_entities(Account.tag) .distinct() .order_by(Account.tag.asc()) .all() ) return [tag for (tag,) in rows if tag] @router.get("/duplicates") def find_duplicate_cookies( db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:view")), ): """检测 CK 管理中同一账号出现多次的记录,不返回 Cookie 敏感字段。""" query = ( db.query( LoginTask.id.label("cookie_id"), LoginTask.account_id.label("account_id"), LoginTask.batch_id.label("batch_id"), LoginTask.finished_at.label("finished_at"), Account.username.label("username"), ) .join(Account, LoginTask.account_id == Account.id) .filter(LoginTask.status == "success") ) # 与 CK 列表保持相同的数据可见范围,客服只能检测自己被分配的账号。 if not user_has_permission(current, "login:view_all"): query = query.filter(Account.assigned_to == current.id) rows = query.order_by( Account.username.asc(), LoginTask.finished_at.desc(), LoginTask.id.desc() ).all() grouped: dict[str, dict] = {} for row in rows: username = (row.username or "").strip() if not username: continue account_key = username.casefold() group = grouped.setdefault( account_key, { "account_key": account_key, "account_names": set(), "account_ids": set(), "cookie_ids": [], "records": [], }, ) group["account_names"].add(username) group["account_ids"].add(row.account_id) group["cookie_ids"].append(row.cookie_id) group["records"].append( { "id": row.cookie_id, "account_id": row.account_id, "batch_id": row.batch_id, "finished_at": _fmt_dt(row.finished_at), } ) duplicate_groups = [] for group in grouped.values(): if len(group["cookie_ids"]) < 2: continue duplicate_groups.append( { "account_key": group["account_key"], "account_names": sorted(group["account_names"]), "cookie_count": len(group["cookie_ids"]), "account_count": len(group["account_ids"]), "cookie_ids": group["cookie_ids"], "account_ids": sorted(group["account_ids"]), "records": group["records"], } ) duplicate_groups.sort(key=lambda item: (-item["cookie_count"], item["account_key"])) return { "success": True, "duplicate_groups": len(duplicate_groups), "duplicate_rows": sum(item["cookie_count"] for item in duplicate_groups), "groups": duplicate_groups, } @router.get("/export") def export_cookies( format: str = "csv", account_names: str = Query(""), db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:export")), ): """导出 Cookie,支持 csv 和 custom 格式。 csv: 账号, Cookie, 时间 custom: 账号----密码----ck """ query = _visible_cookie_tasks_query(db, current) selected_names = _parse_account_names(account_names) if selected_names: if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) query = query.filter(Account.username.in_(selected_names)) tasks = _order_cookie_tasks(query, selected_names).all() # 批量查账号,避免 N+1 account_ids = [t.account_id for t in tasks] accounts_map = {} if account_ids: accs = db.query(Account).filter(Account.id.in_(account_ids)).all() accounts_map = {a.id: a for a in accs} if format == "custom": # 账号----密码----ck lines = [] for t in tasks: acc = accounts_map.get(t.account_id) username = acc.username if acc else "" password = acc.password if acc else "" ck = t.cookie or "" lines.append(f"{username}----{password}----{ck}") content = "\r\n".join(lines) filename = "cookies_custom.txt" media = "text/plain" else: # CSV: 账号, Cookie, 时间 output = io.StringIO() writer = csv.writer(output) writer.writerow(["账号", "Cookie", "时间"]) for t in tasks: acc = accounts_map.get(t.account_id) username = acc.username if acc else "" writer.writerow([username, t.cookie or "", _fmt_dt(t.finished_at) or ""]) content = output.getvalue() filename = "cookies.csv" media = "text/csv" return StreamingResponse( iter([content]), media_type=media, headers={"Content-Disposition": f"attachment; filename={filename}"}, ) def _check_cookies(ids: str, db: Session, current: User, *, detailed: bool) -> dict: """执行 CK 检测;客服操作页仅返回状态,不返回账号衍生信息。""" if not ids: raise HTTPException(status_code=400, detail="请指定记录ID") id_list = [int(x) for x in ids.split(",") if x.strip().isdigit()] if not id_list: raise HTTPException(status_code=400, detail="无效的ID") base_query = ( _visible_cookie_tasks_query(db, current) if detailed else _visible_cookie_operation_tasks_query(db, current) ) tasks = ( base_query.filter(LoginTask.id.in_(id_list)) .filter(LoginTask.status.in_(("success", "relogin_failed"))) .all() ) if not tasks: raise HTTPException(status_code=404, detail="记录不存在") results: list[dict] = [] with ThreadPoolExecutor(max_workers=8) as pool: futures = {pool.submit(_check_one_cookie, task): task for task in tasks} for future in as_completed(futures): try: results.append(future.result()) except Exception as exc: # noqa: BLE001 外部接口与任务边界需要保留宽泛异常兜底 task = futures[future] results.append( { "id": task.id, "valid": False, "message": f"检测异常: {exc}", "fish_ball": None, "nickname": None, "level": None, "checked_at": datetime.now(UTC).isoformat(), } ) results.sort(key=lambda item: item["id"]) # 持久化检测结果,刷新/翻页不丢失 tasks_by_id = {task.id: task for task in tasks} for item in results: task = tasks_by_id.get(item["id"]) if not task: continue task.ck_check_status = "valid" if item["valid"] else "invalid" task.ck_check_result = { "fish_ball": item.get("fish_ball"), "nickname": item.get("nickname"), "level": item.get("level"), "message": item.get("message", ""), } task.ck_checked_at = datetime.now(UTC) db.add( AuditLog( user_id=current.id, username=current.username, action="cookie:check", target=f"检测 {len(results)} 条已分配账号 CK", ) ) db.commit() if not detailed: results = [ {"id": item["id"], "valid": item["valid"], "checked_at": item["checked_at"]} for item in results ] return {"results": results, "success": True} @router.post("/check") def check_cookies( ids: str = "", db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """批量检测斗鱼 Cookie 有效性,供 Cookie 管理页查看详细结果。""" if not user_has_permission(current, "cookie:view"): raise HTTPException(status_code=403, detail="无权限: cookie:view") return _check_cookies(ids, db, current, detailed=True) @router.post("/operations/check") def check_cookie_operations( ids: str = "", db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """脱敏 CK 检测,只返回有效性与检测时间。""" _require_cookie_operation_perm(current) return _check_cookies(ids, db, current, detailed=False) def _start_relogin_tasks( tasks: list[LoginTask], db: Session, current: User, *, action: str = "cookie:relogin", ): """启动重登批次:旧 Cookie 保留到新登录成功后才替换。""" if not tasks: raise HTTPException(status_code=404, detail="记录不存在") task_ids: list[int] = [] skipped: list[int] = [] for task in tasks: acc = db.query(Account).filter(Account.id == task.account_id).first() if not acc: skipped.append(task.id) continue if not acc.password or not acc.email: task.status = "relogin_failed" task.message = "重新登录失败: 账号缺少密码或邮箱,旧 Cookie 已保留" task.finished_at = datetime.now(UTC) skipped.append(task.id) continue if task.status in {"relogin_pending", "relogin_running"}: skipped.append(task.id) continue task.status = "relogin_pending" task.message = "等待重新登录,旧 Cookie 已保留" task.finished_at = None task_ids.append(task.id) if not task_ids: db.commit() raise HTTPException( status_code=400, detail="没有可重登的账号,请检查账号凭据或当前重登状态" ) proxy = db.query(ProxyConfigModel).first() thread_db = SessionLocal() max_login_retries, max_total_time = get_relogin_limits() runner = LoginBatchRunner( db=thread_db, account_ids=[], created_by=current.id, creator_permissions=get_user_permissions(current), max_login_retries=max_login_retries, max_total_time=max_total_time, proxy_config=proxy, log_queue=None, loop=None, concurrency=3, mode="relogin", relogin_task_ids=task_ids, ) batch_id = runner.batch_id db.add( AuditLog( user_id=current.id, username=current.username, action=action, target=f"重登 {len(task_ids)} 条账号 CK", detail="旧 Cookie 将在新登录成功后替换", ) ) db.commit() cookie_relogin_registry.register(batch_id, None, None, runner, owner_id=current.id) def run_relogin_batch(): try: runner.run() except Exception as exc: # noqa: BLE001 外部接口与任务边界需要保留宽泛异常兜底 runner._push_log("error", f"批量重新登录异常: {exc}") message = f"重新登录异常: {exc}(旧 Cookie 已保留)" for task_id in task_ids: runner._mark_relogin_stopped(task_id, message) runner.db.close() finally: cookie_relogin_registry.pop(batch_id) threading.Thread(target=run_relogin_batch, daemon=True).start() return { "batch_id": batch_id, "count": len(task_ids), "skipped": len(skipped), "success": True, } @router.post("/operations/relogin/{batch_id}/stop") def stop_cookie_relogin( batch_id: str, db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """停止当前用户发起的 CK 重登批次。""" _require_cookie_operation_perm(current) batch = cookie_relogin_registry.get(batch_id) if not batch: raise HTTPException( status_code=404, detail="重登批次不存在、已结束或服务已重启" ) if batch.get("owner_id") != current.id and not user_has_permission( current, "login:view_all" ): raise HTTPException(status_code=403, detail="无权限停止该重登批次") batch["runner"].stop() db.add( AuditLog( user_id=current.id, username=current.username, action="cookie:relogin_stop", target=f"停止 CK 重登批次 {batch_id}", ) ) db.commit() return { "message": "已发送停止信号,正在登录的账号会在当前请求结束后停止", "success": True, } def _start_relogin(req: CookieReloginRequest, db: Session, current: User): """按选中的 Cookie 记录启动重新登录。""" if not req.ids: raise HTTPException(status_code=400, detail="请选择要重新登录的 Cookie") tasks = ( _visible_cookie_operation_tasks_query(db, current) .filter(LoginTask.id.in_(req.ids)) .filter(LoginTask.status.in_(("success", "relogin_failed"))) .all() ) return _start_relogin_tasks(tasks, db, current) @router.post("/relogin") def relogin_cookies( req: CookieReloginRequest, db: Session = Depends(get_db), current: User = Depends(require_permission("login:batch")), ): """Cookie 管理页的重登入口,保留既有 login:batch 权限。""" return _start_relogin(req, db, current) @router.post("/operations/relogin") def relogin_cookie_operations( req: CookieReloginRequest, db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """脱敏 CK 操作页的重登入口。""" _require_cookie_operation_perm(current) return _start_relogin(req, db, current) @router.post("/operations/relogin-invalid") def relogin_invalid_cookie_operations( search: str = Query(""), tag: str = Query(""), db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """批量重新登录当前用户可见且最近检测为无效的账号。""" _require_cookie_operation_perm(current) query = _visible_cookie_operation_tasks_query(db, current).filter( LoginTask.ck_check_status == "invalid", LoginTask.status.in_(("success", "relogin_failed")), ) if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) if tag.strip(): query = query.filter(Account.tag == tag.strip()) if search.strip(): pattern = f"%{search.strip()}%" query = query.filter( or_(Account.username.ilike(pattern), Account.tag.ilike(pattern)) ) tasks = query.order_by(LoginTask.id.asc()).all() return _start_relogin_tasks(tasks, db, current, action="cookie:relogin_invalid") @router.get("/{task_id}") def get_cookie( task_id: int, db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """获取单条 Cookie 详情,供复制操作按需读取完整敏感字段。""" if not user_has_permission(current, "cookie:view"): raise HTTPException(status_code=403, detail="无权限: cookie:view") task = ( _visible_cookie_tasks_query(db, current).filter(LoginTask.id == task_id).first() ) if not task: raise HTTPException(status_code=404, detail="记录不存在") acc = db.query(Account).filter(Account.id == task.account_id).first() item = { "id": task.id, "batch_id": task.batch_id, "account_id": task.account_id, "account_username": acc.username if acc else "", "assigned_to": acc.assigned_to if acc else None, "assigned_username": acc.assigned_user.username if acc and acc.assigned_user else None, "created_at": _fmt_dt(task.finished_at), "ck_check_status": task.ck_check_status or "", "ck_check_result": task.ck_check_result, "ck_checked_at": _fmt_dt(task.ck_checked_at), "cookie": "", "cookie_preview": "***", "account_password": "", } if user_has_permission(current, "cookie:view"): cookie = task.cookie or "" item["cookie"] = cookie item["cookie_preview"] = cookie[:50] + "..." if len(cookie) > 50 else cookie item["account_password"] = acc.password if acc else "" return item @router.delete("/batch") def delete_cookies_batch( task_ids: str = "", db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:export")), ): """批量删除 Cookie 记录。""" 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") tasks = _visible_cookie_tasks_query(db, current).filter(LoginTask.id.in_(ids)).all() for t in tasks: t.cookie = "" t.status = "failed" t.message = "Cookie已清除" db.commit() return { "message": f"已删除 {len(tasks)} 条", "deleted": len(tasks), "success": True, } @router.delete("/{task_id}") def delete_cookie( task_id: int, db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:export")), ): """删除一条 Cookie 记录。""" task = ( _visible_cookie_tasks_query(db, current).filter(LoginTask.id == task_id).first() ) if not task: raise HTTPException(status_code=404, detail="记录不存在") task.cookie = "" task.status = "failed" task.message = "Cookie已清除" db.commit() return {"message": "已删除", "success": True}