live-hub-py/web/backend/routers/cookies.py

"""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管理"])


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(
    db: Session = Depends(get_db),
    current: User = Depends(get_current_user),
):
    """查看登录成功的 Cookie 列表。"""
    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}

    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 = _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 = "\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}