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

"""登录任务路由 + WebSocket 实时日志"""

import asyncio
import threading

from fastapi import (
    APIRouter,
    Depends,
    HTTPException,
    WebSocket,
    WebSocketDisconnect,
)
from sqlalchemy import and_, func, or_
from sqlalchemy.orm import Session, load_only

from ..database import SessionLocal, get_db
from ..deps import authenticate_websocket, get_current_user, require_permission
from ..models import Account, LoginTask, User
from ..models import ProxyConfig as ProxyConfigModel
from ..permissions import get_user_permissions, user_has_permission
from ..schemas import LoginBatchRequest, LoginTaskOut
from ..services.audit_service import record_audit
from ..services.login_service import LoginBatchRunner, batch_registry

router = APIRouter(prefix="/api/login", tags=["登录任务"])


def _effective_attempt_status():
    """页面展示用的状态：优先最近一次登录尝试，没有尝试时回退到 CK 状态。

    ``attempt_*`` 才是「这次登录怎么样了」，``status`` 只表示这条记录是否
    持有可用 CK —— 两者分开，停止/失败/重启都不会让账号从列表里消失。
    """
    return func.coalesce(
        func.nullif(LoginTask.attempt_status, ""),
        LoginTask.status,
    )


def _login_task_summary(query):
    """按状态汇总登录任务，避免概览页拉完整任务列表。"""
    rows = (
        query.order_by(None)
        .with_entities(_effective_attempt_status(), func.count(LoginTask.id))
        .group_by(_effective_attempt_status())
        .all()
    )
    status_counts = {status or "": count for status, count in rows}
    return {
        "total": sum(status_counts.values()),
        "success": status_counts.get("success", 0),
        "failed": sum(status_counts.get(status, 0) for status in ("failed", "error")),
        "status_counts": status_counts,
    }


@router.post("/batch")
async def create_batch(
    req: LoginBatchRequest,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("login:batch")),
):
    """创建批量登录或账号检测任务。"""
    if not req.account_ids:
        raise HTTPException(status_code=400, detail="请选择账号")

    # 读取代理配置
    proxy = db.query(ProxyConfigModel).first()

    # 权限过滤账号
    valid_ids = []
    seen_ids = set()
    for aid in req.account_ids:
        if aid in seen_ids:
            continue
        seen_ids.add(aid)
        acc = db.query(Account).filter(Account.id == aid).first()
        if not acc:
            continue
        if (
            not user_has_permission(current, "login:view_all")
            and acc.assigned_to != current.id
        ):
            continue
        valid_ids.append(aid)

    if not valid_ids:
        action_name = "检测" if req.mode == "check" else "登录"
        raise HTTPException(status_code=403, detail=f"没有可{action_name}的账号")

    # 在主事件循环中创建 log_queue，传给后台线程
    log_queue = asyncio.Queue()
    loop = asyncio.get_running_loop()

    # 创建执行器（用独立的 DB 会话，因为在线程中运行）
    thread_db = SessionLocal()
    runner = LoginBatchRunner(
        db=thread_db,
        account_ids=valid_ids,
        created_by=current.id,
        creator_permissions=get_user_permissions(current),
        max_login_retries=req.max_login_retries,
        max_total_time=req.max_total_time,
        proxy_config=proxy,
        log_queue=log_queue,
        loop=loop,
        concurrency=req.concurrency,
        api_strategy=req.api_strategy,
        mode=req.mode,
    )

    batch_id = runner.batch_id

    # 先注册到批次注册表，再启动线程，确保 WebSocket 连接时能找到
    batch_registry.register(batch_id, log_queue, loop, runner)

    # 启动线程
    thread = threading.Thread(target=runner.run, daemon=True)
    thread.start()

    # 审计：批次是谁在什么时候发起、多少个账号、什么模式。
    # 自动重登与人工登录共用同一条链路，没有这条记录就无法追溯来源。
    try:
        record_audit(
            db,
            current,
            action="login:batch",
            target=batch_id,
            detail={
                "mode": req.mode,
                "requested": len(req.account_ids),
                "accepted": len(valid_ids),
                "concurrency": req.concurrency,
                "api_strategy": req.api_strategy,
            },
        )
        db.commit()
    except Exception:  # noqa: BLE001 审计失败不应影响已启动的批次
        db.rollback()

    return {"batch_id": batch_id, "count": len(valid_ids), "success": True}


def _visible_tasks_query(db: Session, current: User, batch_id: str | None):
    """按权限与批次过滤登录任务。客服只能看自己名下账号的任务。

    这里已经**不需要 defer cookie** —— 登录记录不再持有凭证（2026-10-04 CK 解耦）。
    """
    query = db.query(LoginTask).join(Account, LoginTask.account_id == Account.id)

    if not user_has_permission(current, "login:view_all"):
        query = query.filter(Account.assigned_to == current.id)

    if batch_id:
        # 记录持有的 CK 属于产出它的批次，最近一次尝试属于新批次，两者都要算。
        query = query.filter(
            or_(LoginTask.batch_id == batch_id, LoginTask.attempt_batch_id == batch_id)
        )
    return query


def _serialize_tasks(db: Session, rows) -> list[LoginTaskOut]:
    """把任务行拼上账号名后序列化。"""
    account_ids = [t.account_id for t in rows]
    accounts_map = {}
    if account_ids:
        accs = (
            db.query(Account)
            .options(load_only(Account.id, Account.username))
            .filter(Account.id.in_(account_ids))
            .all()
        )
        accounts_map = {a.id: a.username for a in accs}

    result = []
    for t in rows:
        attempt_status = t.attempt_status or ""
        result.append(
            LoginTaskOut(
                id=t.id,
                batch_id=t.attempt_batch_id or t.batch_id,
                account_id=t.account_id,
                account_username=accounts_map.get(t.account_id, ""),
                status=attempt_status or t.status,
                cookie="",
                message=(t.attempt_message if attempt_status else "")
                or t.message
                or "",
                created_by=t.created_by,
                created_at=t.created_at,
                finished_at=t.attempt_at or t.finished_at,
            )
        )
    return result


@router.get("/tasks")
def list_tasks(
    batch_id: str | None = None,
    # 默认值保持 None（而不是 Query(None)）：这个接口会被测试与内部代码直接
    # 调用，用 Query() 做默认值时拿到的会是那个 Query 对象而非 None。
    page: int | None = None,
    page_size: int = 50,
    status: str | None = None,
    db: Session = Depends(get_db),
    current: User = Depends(get_current_user),
):
    """查看登录任务列表。

    **不传 `page` 时返回裸数组**（上限 200），沿用老调用方（Cookie 页等）的
    契约；传了 `page` 才返回 ``{items, total, page, page_size}``，供表格分页。

    传 `page` 时同时返回 `total`，是为了让前端能「全选所有筛选结果」——
    否则跨页选择无从下手，而任务量早就过了单页能装下的规模。
    """
    query = _visible_tasks_query(db, current, batch_id)
    if status:
        # 前端看到的是「最近一次尝试」的状态，与列表序列化口径保持一致。
        query = query.filter(
            or_(
                and_(
                    LoginTask.attempt_status != "", LoginTask.attempt_status == status
                ),
                and_(LoginTask.attempt_status == "", LoginTask.status == status),
            )
        )

    ordered = query.order_by(LoginTask.id.desc())

    if page is None:
        return _serialize_tasks(db, ordered.limit(200).all())

    page = max(1, int(page))
    page_size = max(1, min(int(page_size), 500))
    total = ordered.count()
    rows = ordered.offset((page - 1) * page_size).limit(page_size).all()
    return {
        "items": _serialize_tasks(db, rows),
        "total": total,
        "page": page,
        "page_size": page_size,
    }


@router.get("/tasks/summary")
def tasks_summary(
    db: Session = Depends(get_db),
    current: User = Depends(get_current_user),
):
    """登录任务统计。"""
    query = db.query(LoginTask).join(Account, LoginTask.account_id == Account.id)
    if not user_has_permission(current, "login:view_all"):
        query = query.filter(Account.assigned_to == current.id)
    return _login_task_summary(query)


def _in_flight_clause():
    """不可删除的记录：正被批次处理的**在途尝试**。

    ⚠️ **CK 自 2026-10-04 起不再是护栏理由** —— CK 已迁到 `accounts.cookie`，
    删登录历史不再影响任何账号的登录态（这正是解耦的目的）。但「在途」必须保留：
    删掉会让批次线程找不到记录；停批次请走 ``/stop/{batch_id}``。
    """
    in_flight = ("pending", "running")
    return or_(
        and_(
            LoginTask.attempt_status != "",
            LoginTask.attempt_status.in_(in_flight),
        ),
        and_(LoginTask.attempt_status == "", LoginTask.status.in_(in_flight)),
    )


def _deletable_clause():
    """可安全删除的记录：不在途。"""
    return ~_in_flight_clause()


def _is_in_flight(task: LoginTask) -> bool:
    """单条记录是否在途（与 `_in_flight_clause` 同口径）。"""
    in_flight = ("pending", "running")
    if task.attempt_status:
        return task.attempt_status in in_flight
    return (task.status or "") in in_flight


@router.delete("/tasks/{task_id}")
def delete_task(
    task_id: int,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("login:batch")),
):
    """删除单个登录任务（**在途的记录不允许删除**；删历史不影响 CK）。"""
    task = db.query(LoginTask).filter(LoginTask.id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")
    if _is_in_flight(task):
        raise HTTPException(
            status_code=409,
            detail="该记录正在被批次处理，删除会让批次失去目标；请先停止批次",
        )
    db.delete(task)
    db.commit()
    return {"message": "已删除", "success": True}


@router.delete("/tasks")
def delete_tasks(
    task_ids: str = "",
    batch_id: str | None = None,
    status: str | None = None,
    all_matching: bool = False,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("login:batch")),
):
    """批量删除登录任务。三种模式，互斥：

    1. ``task_ids``：删指定 id（表格勾选）；
    2. ``all_matching=true``：删当前筛选条件下的**全部**（跨页全选用）；
    3. 两者都没给 → 400，绝不默认全删。

    两条护栏（都会在响应里报告跳过数量）：

    - **在途（pending/running）的记录不删**：删掉会让批次线程找不到记录；
      停批次请走 ``/stop/{batch_id}``。
    - ~~持有 CK 的记录不删~~ —— **这条护栏已随 CK 解耦（2026-10-04）取消**：
      CK 现在挂在账号上，删登录历史不影响登录态。历史终于可以自由清理。
    """
    # 护栏：不删在途任务，避免批次线程与数据库状态打架。
    not_in_flight = _deletable_clause()

    if task_ids:
        ids = [int(x) for x in task_ids.split(",") if x.strip().isdigit()]
        if not ids:
            raise HTTPException(status_code=400, detail="无效的任务ID")
    elif all_matching:
        # 带 join 的 query 不能直接 bulk delete（SQLAlchemy 会拒绝），
        # 所以先按权限与筛选条件查出 id，再按主键删。
        query = _visible_tasks_query(db, current, batch_id)
        if status:
            query = query.filter(
                or_(
                    and_(
                        LoginTask.attempt_status != "",
                        LoginTask.attempt_status == status,
                    ),
                    and_(LoginTask.attempt_status == "", LoginTask.status == status),
                )
            )
        ids = [row[0] for row in query.with_entities(LoginTask.id).all()]
    else:
        raise HTTPException(status_code=400, detail="请指定任务ID或 all_matching=true")

    if not ids:
        return {
            "message": "没有符合条件的任务",
            "deleted": 0,
            "skipped": 0,
            "success": True,
        }

    rows = (
        db.query(LoginTask.id, not_in_flight.label("ok_to_delete"))
        .filter(LoginTask.id.in_(ids))
        .all()
    )
    target_ids = [r[0] for r in rows if r[1]]
    skipped = len(rows) - len(target_ids)

    if not target_ids:
        return {
            "message": f"没有可删除的记录（{skipped} 条在途，已跳过）",
            "deleted": 0,
            "skipped": skipped,
            "success": True,
        }

    deleted = (
        db.query(LoginTask)
        .filter(LoginTask.id.in_(target_ids))
        .delete(synchronize_session=False)
    )
    db.commit()

    message = f"已删除 {deleted} 个任务"
    if skipped:
        message += f"（{skipped} 条在途，已跳过）"
    return {"message": message, "deleted": deleted, "skipped": skipped, "success": True}


@router.post("/stop/{batch_id}")
def stop_batch(
    batch_id: str,
    current: User = Depends(require_permission("login:batch")),
):
    batch = batch_registry.get(batch_id)
    if batch:
        if batch.get("finished"):
            # 已结束但 WS 尚未清理（或从未有 WS 订阅）：回收注册表条目。
            batch_registry.pop(batch_id)
            raise HTTPException(status_code=404, detail="批次已结束")
        batch["runner"].stop()
        return {"message": "已发送停止信号", "success": True}
    raise HTTPException(status_code=404, detail="批次不存在或已结束")


@router.websocket("/ws/login/{batch_id}")
async def ws_login_logs(websocket: WebSocket, batch_id: str):
    """WebSocket 推送登录实时日志（需认证）。"""
    # 认证：从 cookie 或 token query param 验证用户身份
    user = authenticate_websocket(websocket)
    if not user:
        await websocket.close(code=1008, reason="未授权")
        return

    # 归属校验：非 view_all 只能订阅自己创建的批次（与虎牙批次一致）。
    if not user_has_permission(user, "login:view_all"):
        check_db = SessionLocal()
        try:
            owned = (
                check_db.query(LoginTask.id)
                .filter(LoginTask.batch_id == batch_id, LoginTask.created_by == user.id)
                .first()
            )
        finally:
            check_db.close()
        if not owned:
            await websocket.close(code=1008, reason="无权访问该任务批次")
            return

    await websocket.accept()

    # 从已注册的批次中获取 log_queue（由 create_batch 创建）
    batch = batch_registry.get(batch_id)
    if not batch:
        await websocket.send_json({"level": "error", "message": "批次不存在或已结束"})
        await websocket.close()
        return

    log_queue: asyncio.Queue = batch["log_queue"]

    try:
        while True:
            try:
                msg = await asyncio.wait_for(log_queue.get(), timeout=30)
                await websocket.send_json(msg)
                # 收到 result 表示任务结束
                if msg.get("level") == "result":
                    await asyncio.sleep(0.1)
                    break
            except TimeoutError:
                await websocket.send_json({"level": "heartbeat", "message": ""})
    except WebSocketDisconnect:
        pass
    finally:
        # 只在批次真正结束后清理注册表，客户端断线/刷新页面后重连仍可继续订阅。
        latest = batch_registry.get(batch_id)
        if latest and latest.get("finished"):
            batch_registry.pop(batch_id)