"""Cookie 管理路由""" from datetime import timezone 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管理"]) 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), } # 分页列表默认只返回预览,复制/导出时再获取完整敏感字段。 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("/{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), "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.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.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}