NeteaseUrl/web/billing.py

# -*- coding: utf-8 -*-
"""计费依赖（账本预扣/退款），口径与登录闸一致。

历史实现用请求上下文（flask.g）暂存请求级账本键；本实现无 g，改用
「依赖返回值携带 ledger_key + 路由内显式退款」的方式，语义完全一致：
- 预扣与开账同事务：auth_system.reserve_and_open
- 失败退款：只退「仍 open 的账本行」（与装饰器 _refund_open 同口径，杜绝双退）
- 任务接手：auth_system.ledger_rename 把请求级行改挂 task_id（worker 结算唯一入口）

金额相关逻辑一律走 auth_system 既有函数，本模块不新造账务语义。
"""
import logging
import secrets
from collections.abc import AsyncGenerator, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass

import auth_system
from web.deps import SiteUser
from web.errors import AppError

logger = logging.getLogger('web.billing')


@dataclass
class Billing:
    """一次请求的计费上下文。"""
    user: SiteUser
    ledger_key: str | None      # 请求级账本键（unlimited 用户为 None）
    need: int                   # 本次预扣次数（free 路径为 0）
    free: bool

    async def refund_open(self, reason: str = 'rejected') -> None:
        """仅对仍 open 的账本行全额退款（被 rename/结算过则跳过）。"""
        key = self.ledger_key
        if not key:
            return
        try:
            row = await auth_system.ledger_get(key)
            if row and row['status'] == 'open':
                await auth_system.ledger_settle(key, 0, reason)
        except Exception:
            pass

    async def settle(self, delivered: int, status: str = 'delivered') -> None:
        """交付结算（json 模式等无 task_id 的场景）。"""
        key = self.ledger_key
        if not key:
            return
        try:
            await auth_system.ledger_settle(key, delivered, status)
        except Exception as e:
            logger.error('账本结算失败（重启对账兜底）: %s %s', key, e)

    async def rename_to(self, task_id: str) -> bool:
        """把请求级账本行改挂 task_id（任务接手账目）。无请求级行时开 reserved=0 的账。"""
        if not self.ledger_key:
            await auth_system.ledger_open(task_id, self.user.id, 0)
            return True
        return await auth_system.ledger_rename(self.ledger_key, task_id)

    @asynccontextmanager
    async def scope(self) -> AsyncGenerator['Billing', None]:
        """计费作用域：退出时自动退还仍 open 的账本行。

        把「任何失败路径必须退款」这条不变量从路由的每条早退/except 里
        收敛到这里。正确性依赖 refund_open 的既有幂等口径（只退仍 open
        的行），因此：
        - 错误早退 / 异常转错误响应 → 正常退出作用域 → 自动退；
        - settle()/rename_to() 已消费的行不在退款范围，成功路径自然豁免，
          无需手动「取消自动退」；
        - 未捕获异常先退款再原样穿透，由路由既有 except 转错误响应；
        - 客户端断连（CancelledError）与迁移前行为一致：退款可能被一并
          取消，open 行留给重启对账兜底，不额外加 shield 复杂化。
        """
        try:
            yield self
        finally:
            await self.refund_open('rejected')


async def require_billing(
    user: SiteUser | None,
    *,
    action: str = '',
    free: bool = False,
    cost: Callable[[], int] | None = None,
    ip: str = '-',
) -> Billing:
    """登录 + 预扣闸（对应 Flask login_required）。

    free=True：仅要求登录，不扣次数（解析/试听类）。
    cost: 预扣次数计算（下载类：单曲 1 / 批量按去重后歌曲数）。
    """
    if user is None:
        raise AppError('请先登录后再使用解析功能', 401, error_code='unauthorized')

    fresh = await auth_system.get_user_by_id(user.id)
    if fresh is None:
        raise AppError('登录已失效，请重新登录', 401, error_code='unauthorized')

    credits = int(fresh['duration_seconds'] or 0)
    unlimited = fresh['role'] == 'admin' or credits == -1

    if free or unlimited:
        return Billing(user=user, ledger_key=None, need=0, free=free)

    need = 1
    if cost is not None:
        try:
            need = max(1, int(cost()))
        except Exception:
            need = 1

    key = 'req-' + secrets.token_hex(8)
    remaining = await auth_system.reserve_and_open(fresh['id'], need, key)
    if remaining is None:
        await auth_system.add_log('fail', fresh['username'], ip,
                                  action or '下载', f'剩余次数不足（本次需 {need} 次）')
        raise AppError(f'下载次数已用尽（本次需 {need} 次，当前 {credits} 次），'
                       '请充值卡密后继续下载', 401, error_code='expired')
    return Billing(user=user, ledger_key=key, need=need, free=False)