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

"""斗鱼账号检测路由。"""

from __future__ import annotations

import threading

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from ..database import get_db
from ..deps import get_current_user
from ..models import Account, User
from ..models import ProxyConfig as ProxyConfigModel
from ..permissions import user_has_permission
from ..schemas import AccountCheckBatchOut, AccountCheckBatchRequest
from ..services.account_check_service import (
    account_check_registry,
    build_inputs_from_accounts,
    parse_account_check_lines,
)

router = APIRouter(prefix="/api/account-check", tags=["账号检测"])


def _get_runner_or_404(batch_id: str):
    """读取账号检测批次。"""
    runner = account_check_registry.get(batch_id)
    if not runner:
        raise HTTPException(status_code=404, detail="批次不存在或服务已重启")
    return runner


def _require_account_check_perm(user: User) -> None:
    """账号检测权限；兼容已有批量登录权限。"""
    if not (
        user_has_permission(user, "account:check")
        or user_has_permission(user, "login:batch")
    ):
        raise HTTPException(status_code=403, detail="权限不足")


@router.post("/batches", response_model=AccountCheckBatchOut)
def create_account_check_batch(
    req: AccountCheckBatchRequest,
    db: Session = Depends(get_db),
    current: User = Depends(get_current_user),
):
    """启动斗鱼账号检测批次。

    两种来源**二选一**：
    - `text`：手动粘贴（旧模式，逻辑一字未动）；
    - `account_ids`：按账号选择（新模式）—— 凭据从库里取，结果写回账号时间线。
    """
    _require_account_check_perm(current)
    has_text = bool((req.text or "").strip())
    has_ids = bool(req.account_ids)
    has_filter = bool(req.select_all_matching)
    picked = has_ids or has_filter
    if has_text and picked:
        raise HTTPException(status_code=400, detail="手动输入与选择账号只能二选一")
    if not has_text and not picked:
        raise HTTPException(status_code=400, detail="请粘贴账号或选择要检测的账号")

    if picked:
        # 可见性沿用账号域的口径：客服只能检测自己名下的账号。
        from .accounts import _filter_accounts_query, _visible_accounts_query

        query = _visible_accounts_query(db, current)
        if has_filter:
            # 「按筛选全选」：与账号列表同一套筛选条件（标签/状态/CK/绑定/搜索），
            # 所以列表上筛出什么，这里就检测什么，不会两处口径不一致。
            query = _filter_accounts_query(
                db,
                query,
                current,
                search=req.search,
                tag=req.tag,
                status=req.status,
                ck_status=req.ck_status,
                bind_status=req.bind_status,
            )
            rows = query.order_by(Account.id.asc()).all()
            ordered_ids = [int(a.id) for a in rows]
        else:
            ordered_ids = [int(i) for i in req.account_ids]
            rows = query.filter(Account.id.in_(ordered_ids)).all()
            found = {int(a.id) for a in rows}
            missing = [i for i in ordered_ids if i not in found]
            if missing:
                raise HTTPException(
                    status_code=403,
                    detail=f"以下账号不存在或无权查看：{missing[:10]}",
                )
        if not rows:
            raise HTTPException(status_code=400, detail="当前条件下没有可检测的账号")
        accounts = build_inputs_from_accounts(rows, ordered_ids)
    else:
        try:
            accounts = parse_account_check_lines(req.text)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    if not accounts:
        raise HTTPException(status_code=400, detail="没有识别到有效账号")

    proxy_config = db.query(ProxyConfigModel).first() if req.use_proxy else None
    runner = account_check_registry.create(
        accounts=accounts,
        created_by=current.id,
        concurrency=req.concurrency,
        max_login_retries=req.max_login_retries,
        max_total_time=req.max_total_time,
        use_proxy=req.use_proxy,
        proxy_config=proxy_config,
    )
    thread = threading.Thread(target=runner.run, daemon=True)
    thread.start()
    return runner.snapshot()


@router.get("/batches/{batch_id}", response_model=AccountCheckBatchOut)
def get_account_check_batch(
    batch_id: str,
    current: User = Depends(get_current_user),
):
    """查询斗鱼账号检测批次。"""
    _require_account_check_perm(current)
    runner = _get_runner_or_404(batch_id)
    return runner.snapshot()


@router.post("/batches/{batch_id}/stop")
def stop_account_check_batch(
    batch_id: str,
    current: User = Depends(get_current_user),
):
    """停止斗鱼账号检测批次。"""
    _require_account_check_perm(current)
    runner = _get_runner_or_404(batch_id)
    runner.stop()
    return {"message": "已发送停止信号", "success": True}


@router.get("/batches/{batch_id}/download")
def download_account_check_batch(
    batch_id: str,
    current: User = Depends(get_current_user),
):
    """下载斗鱼账号检测分类结果 zip。"""
    _require_account_check_perm(current)
    runner = _get_runner_or_404(batch_id)
    snapshot = runner.snapshot()
    if snapshot["status"] in {"pending", "running"}:
        raise HTTPException(status_code=400, detail="批次尚未完成")

    content, filename = runner.build_zip()
    return StreamingResponse(
        iter([content]),
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )