yml2213
a0db26def3
任务量早就过了单页能装下的规模（库里 1.4 万条，单批就 177 条），但接口
硬编码 limit(200)、前端 pagination={false}，于是「全选」其实只能选到前
200 条，失败任务清不干净。
- 接口加 page/page_size/status：不传 page 时仍返回裸数组（Cookie 页等
老调用方依赖这个契约，不能改）；传了才返回 {items,total}，total 供跨页全选。
- 删除支持 all_matching：按当前筛选条件删全部，不必先翻完所有页去勾。
- 前端：分页控件（20/50/100/200/500，记住选择）、状态筛选、表头
「选择全部 N 条（含其他页）」、「清空失败」一键操作。
两条护栏：
- **pending/running 一律不删** —— 删掉会让批次线程找不到记录；要停批次
走 /stop/{batch_id}。
- 不给任何条件就拒绝删除，绝不默认全删。
顺带：带 join 的 query 不能 bulk delete（SQLAlchemy 会拒绝），改成先查
id 再按主键删。
418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""登录任务路由 + WebSocket 实时日志"""
|
||
|
||
import asyncio
|
||
import threading
|
||
|
||
from fastapi import (
|
||
APIRouter,
|
||
Depends,
|
||
HTTPException,
|
||
WebSocket,
|
||
WebSocketDisconnect,
|
||
)
|
||
from sqlalchemy import and_, func, or_
|
||
from sqlalchemy.orm import Session, defer, 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):
|
||
"""按权限与批次过滤登录任务。客服只能看自己名下账号的任务。"""
|
||
query = (
|
||
db.query(LoginTask)
|
||
.options(defer(LoginTask.cookie))
|
||
.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)
|
||
|
||
|
||
@router.delete("/tasks/{task_id}")
|
||
def delete_task(
|
||
task_id: int,
|
||
db: Session = Depends(get_db),
|
||
current: User = Depends(require_permission("login:batch")),
|
||
):
|
||
"""删除单个登录任务。"""
|
||
task = db.query(LoginTask).filter(LoginTask.id == task_id).first()
|
||
if not task:
|
||
raise HTTPException(status_code=404, 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}``。
|
||
"""
|
||
# 护栏：不删在途任务，避免批次线程与数据库状态打架。
|
||
in_flight = ("pending", "running")
|
||
not_in_flight = or_(
|
||
and_(
|
||
LoginTask.attempt_status != "",
|
||
LoginTask.attempt_status.notin_(in_flight),
|
||
),
|
||
and_(LoginTask.attempt_status == "", LoginTask.status.notin_(in_flight)),
|
||
)
|
||
|
||
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, "success": True}
|
||
|
||
target_ids = [
|
||
row[0]
|
||
for row in db.query(LoginTask.id)
|
||
.filter(LoginTask.id.in_(ids))
|
||
.filter(not_in_flight)
|
||
.all()
|
||
]
|
||
if not target_ids:
|
||
return {
|
||
"message": "没有可删除的任务（在途任务已跳过）",
|
||
"deleted": 0,
|
||
"success": True,
|
||
}
|
||
|
||
deleted = (
|
||
db.query(LoginTask)
|
||
.filter(LoginTask.id.in_(target_ids))
|
||
.delete(synchronize_session=False)
|
||
)
|
||
db.commit()
|
||
return {"message": f"已删除 {deleted} 个任务", "deleted": deleted, "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)
|