"""Cookie 管理路由""" from datetime import timezone from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy.orm import Session 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管理"]) @router.get("") def list_cookies( db: Session = Depends(get_db), current: User = Depends(get_current_user), ): """查看登录成功的 Cookie 列表。""" query = db.query(LoginTask).filter(LoginTask.status == "success") # 客服只能看自己账号的 if not user_has_permission(current, "login:view_all"): query = query.join(Account, LoginTask.account_id == Account.id).filter( Account.assigned_to == current.id ) tasks = query.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} 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), } # 只有有 cookie:view 权限才返回 cookie 内容 if user_has_permission(current, "cookie:view"): cookie = t.cookie or "" item["cookie"] = cookie item["cookie_preview"] = cookie[:50] + "..." if len(cookie) > 50 else cookie else: item["cookie"] = "" item["cookie_preview"] = "***" result.append(item) return result @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 = db.query(LoginTask).filter( LoginTask.status == "success" ).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 = "\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("/{task_id}") def delete_cookie( task_id: int, db: Session = Depends(get_db), current: User = Depends(require_permission("cookie:export")), ): """删除一条 Cookie 记录。""" task = db.query(LoginTask).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}