"""Cookie 管理路由""" from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timezone import requests from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse from sqlalchemy import or_ from sqlalchemy.orm import Session, defer, joinedload import io import csv from ..database import get_db from ..models import User, LoginTask, Account from ..deps import get_current_user, require_permission from ..permissions import user_has_permission def _fmt_dt(dt) -> str | None: """将 datetime 格式化为带时区的 ISO 字符串。""" if dt is None: return None if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) return dt.isoformat() router = APIRouter(prefix="/api/cookies", tags=["Cookie管理"]) # 斗鱼 Cookie 有效性检测接口(抓包参考:钱包中心鱼丸余额 + 用户等级详情) CHECK_UA = ( "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 " "(KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36" ) FISH_BALL_API = "https://www.douyu.com/wjapi/nc/exchange/fishBall" USER_LEVEL_API = "https://www.douyu.com/japi/interactnc/web/userLevel/userLevelDetail" def _check_one_cookie(task: LoginTask) -> dict: """检测单条斗鱼 Cookie 有效性:鱼丸 + 用户等级双接口均通过才算有效。 附带鱼丸数、昵称/等级与检测时间;任一接口失败会说明原因。 """ cookie = task.cookie or "" checked_at = datetime.now(timezone.utc).isoformat() base = {"id": task.id, "checked_at": checked_at, "fish_ball": None, "nickname": None, "level": None} if not cookie: return {**base, "valid": False, "message": "Cookie 为空"} fish_ok = False fish_msg = "" fish_ball = None try: fish_data = requests.get( FISH_BALL_API, params={"appCode": "YJTX"}, headers={ "User-Agent": CHECK_UA, "Accept": "application/json, text/plain, */*", "Referer": "https://www.douyu.com/member/walletcenter", "Cookie": cookie, }, timeout=(5, 10), ).json() if isinstance(fish_data, dict) and fish_data.get("error") in (0, "0"): fish_ok = True fish_ball = (fish_data.get("data") or {}).get("count") if isinstance(fish_data.get("data"), dict) else None else: fish_msg = str(fish_data.get("msg") or fish_data.get("error") or "响应异常") if isinstance(fish_data, dict) else "响应异常" except Exception as exc: fish_msg = f"请求失败: {exc}" level_ok = False level_msg = "" nickname = None level = None try: level_data = requests.get( USER_LEVEL_API, params={"rid": "0"}, headers={ "User-Agent": CHECK_UA, "Accept": "application/json, text/plain, */*", "Referer": "https://www.douyu.com/pages/ord-user-level?clientType=web", "Cookie": cookie, }, timeout=(5, 10), ).json() if isinstance(level_data, dict) and level_data.get("error") in (0, "0"): level_ok = True info = level_data.get("data") if isinstance(level_data.get("data"), dict) else {} nickname = str(info.get("nn") or "") or None level = info.get("lv") else: level_msg = str(level_data.get("msg") or level_data.get("error") or "响应异常") if isinstance(level_data, dict) else "响应异常" except Exception as exc: level_msg = f"请求失败: {exc}" valid = fish_ok and level_ok if valid: message = "有效" else: parts = [f"鱼丸接口: {'ok' if fish_ok else (fish_msg or '失败')}", f"等级接口: {'ok' if level_ok else (level_msg or '失败')}"] message = ";".join(parts) return { **base, "valid": valid, "message": message[:200], "fish_ball": fish_ball, "nickname": nickname, "level": level, } 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 @router.get("") def list_cookies( search: 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 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() if search_text: pattern = f"%{search_text}%" if user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id) query = query.outerjoin(User, Account.assigned_to == User.id).filter(or_( Account.username.ilike(pattern), User.username.ilike(pattern), )) total = None if page is not None: total = query.order_by(None).count() query = query.order_by(LoginTask.finished_at.desc()) 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, "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。""" 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.count() assigned_count = query.filter(Account.assigned_to.isnot(None)).count() return { "total": total, "assigned_count": assigned_count, "unassigned_count": max(0, total - assigned_count), } @router.get("/export") def export_cookies( format: str = "csv", db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:export")), ): """导出 Cookie,支持 csv 和 custom 格式。 csv: 账号, Cookie, 时间 custom: 账号----密码----ck """ tasks = _visible_cookie_tasks_query(db, current).order_by(LoginTask.finished_at.desc()).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}"}, ) @router.post("/check") def check_cookies( ids: str = "", db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """批量检测斗鱼 Cookie 有效性(鱼丸接口),返回每条有效性与鱼丸数/昵称。""" 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") tasks = _visible_cookie_tasks_query(db, current).filter(LoginTask.id.in_(id_list)).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: task = futures[future] results.append({ "id": task.id, "valid": False, "message": f"检测异常: {exc}", "fish_ball": None, "nickname": None, "level": None, "checked_at": datetime.now(timezone.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(timezone.utc) db.commit() return {"results": results, "success": True} @router.get("/{task_id}") def get_cookie( task_id: int, db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """获取单条 Cookie 详情,供复制操作按需读取完整敏感字段。""" 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}