shaguabijia-app-server/app/core/limit_policy.py

"""Unified global limits and per-phone/device policy overrides.

The registry is the single source of truth for the whitelist page.  Existing
constants remain as backwards-compatible defaults, while business call sites
resolve an effective value here.
"""
from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from datetime import UTC, datetime

from sqlalchemy import delete, or_, select
from sqlalchemy.orm import Session

from app.core.config_schema import (
    AD_FEED_DAILY_LIMIT_KEY,
    AD_REWARD_VIDEO_DAILY_LIMIT_KEY,
    COMPARE_DAILY_LIMIT_KEY,
    GUIDE_VIDEO_MAX_PLAYS_KEY,
    LIMIT_POLICY_GLOBAL_KEY,
    PHONE_REBIND_DAYS_KEY,
    RISK_COMPARE_DAILY_THRESHOLD_KEY,
    RISK_ONECLICK_DAILY_THRESHOLD_KEY,
    RISK_SMS_HOURLY_THRESHOLD_KEY,
    SMS_CODE_MAX_FAILED_ATTEMPTS_KEY,
    SMS_LOGIN_HOURLY_LIMIT_KEY,
    SMS_PHONE_COOLDOWN_SECONDS_KEY,
    SMS_SEND_DAILY_LIMIT_KEY,
    SMS_SEND_HOURLY_LIMIT_KEY,
    WECHAT_BIND_SMS_HOURLY_LIMIT_KEY,
    WECHAT_CONFLICT_HOURLY_LIMIT_KEY,
)
from app.models.app_config import AppConfig
from app.models.limit_policy import LimitPolicyOverride
from app.models.user import User
from app.repositories import app_config

MODE_INHERIT = "inherit"
MODE_OVERRIDE = "override"
MODE_UNLIMITED = "unlimited"
MODE_SUPPRESS_ALERT = "suppress_alert"

SUBJECT_PHONE = "phone"
SUBJECT_DEVICE = "device"
SUBJECT_TYPES = (SUBJECT_PHONE, SUBJECT_DEVICE)
SUBJECT_PRECEDENCE = {SUBJECT_DEVICE: 0, SUBJECT_PHONE: 1}
LEGACY_IP_DEVICE_PREFIX = "legacy-ip:"


@dataclass(frozen=True)
class RuleDefinition:
    code: str
    label: str
    group: str
    config_key: str
    default_limit: int
    window_label: str
    subject_types: tuple[str, ...]
    min_value: int = 1
    max_value: int = 100_000
    allow_unlimited: bool = True
    supports_reset: bool = True
    alert_only: bool = False
    legacy_config_key: str | None = None
    legacy_json_field: str | None = None

    @property
    def allowed_modes(self) -> tuple[str, ...]:
        if self.alert_only:
            return (MODE_SUPPRESS_ALERT,)
        return (MODE_UNLIMITED,) if self.allow_unlimited else ()


RULES: tuple[RuleDefinition, ...] = (
    RuleDefinition(
        "compare.start.daily",
        "每日发起比价次数",
        "比价",
        COMPARE_DAILY_LIMIT_KEY,
        100,
        "北京时间自然日",
        SUBJECT_TYPES,
    ),
    RuleDefinition(
        "sms.send.hourly",
        "短信每小时成功发送次数",
        "短信与登录",
        SMS_SEND_HOURLY_LIMIT_KEY,
        5,
        "固定 1 小时窗口",
        SUBJECT_TYPES,
        max_value=10_000,
    ),
    RuleDefinition(
        "sms.send.daily",
        "短信 24 小时成功发送次数",
        "短信与登录",
        SMS_SEND_DAILY_LIMIT_KEY,
        20,
        "固定 24 小时窗口",
        SUBJECT_TYPES,
    ),
    RuleDefinition(
        "sms.phone.cooldown",
        "同手机号短信发送冷却",
        "短信与登录",
        SMS_PHONE_COOLDOWN_SECONDS_KEY,
        60,
        "秒",
        (SUBJECT_PHONE,),
        min_value=0,
        max_value=86_400,
    ),
    RuleDefinition(
        "sms.code.failed_attempts",
        "单验证码最大失败次数",
        "短信与登录",
        SMS_CODE_MAX_FAILED_ATTEMPTS_KEY,
        5,
        "单个验证码",
        (SUBJECT_PHONE,),
        max_value=100,
    ),
    RuleDefinition(
        "sms.login.hourly",
        "短信登录每小时尝试次数",
        "短信与登录",
        SMS_LOGIN_HOURLY_LIMIT_KEY,
        5,
        "固定 1 小时窗口",
        SUBJECT_TYPES,
        max_value=10_000,
    ),
    RuleDefinition(
        "wechat.bind.hourly",
        "微信短信绑定每小时尝试次数",
        "短信与登录",
        WECHAT_BIND_SMS_HOURLY_LIMIT_KEY,
        5,
        "固定 1 小时窗口",
        SUBJECT_TYPES,
        max_value=10_000,
    ),
    RuleDefinition(
        "wechat.conflict.hourly",
        "微信冲突处理每小时尝试次数",
        "短信与登录",
        WECHAT_CONFLICT_HOURLY_LIMIT_KEY,
        5,
        "固定 1 小时窗口",
        SUBJECT_TYPES,
        max_value=10_000,
    ),
    RuleDefinition(
        "ad.reward_video.daily",
        "激励视频每日发奖次数",
        "广告",
        AD_REWARD_VIDEO_DAILY_LIMIT_KEY,
        500,
        "北京时间自然日",
        (SUBJECT_PHONE,),
        legacy_config_key="ad_daily_limit",
    ),
    RuleDefinition(
        "ad.feed.daily",
        "Draw 信息流每日发奖次数",
        "广告",
        AD_FEED_DAILY_LIMIT_KEY,
        500,
        "北京时间自然日",
        (SUBJECT_PHONE,),
        legacy_config_key="ad_daily_limit",
    ),
    RuleDefinition(
        "ad.reward_video.cooldown",
        "激励视频发奖后冷却秒数",
        "广告",
        "ad_cooldown_sec",
        3,
        "秒",
        (SUBJECT_PHONE,),
        min_value=0,
        max_value=86_400,
    ),
    RuleDefinition(
        "guide.video.lifetime",
        "领券引导视频最大播放次数",
        "引导与账号",
        GUIDE_VIDEO_MAX_PLAYS_KEY,
        3,
        "账号生命周期",
        (SUBJECT_PHONE,),
        min_value=0,
        max_value=50,
        legacy_config_key="coupon_guide_video",
        legacy_json_field="max_plays",
    ),
    RuleDefinition(
        "phone.rebind.days",
        "手机/微信换绑冷却天数",
        "引导与账号",
        PHONE_REBIND_DAYS_KEY,
        30,
        "自然日",
        (SUBJECT_PHONE,),
        min_value=0,
        max_value=3650,
    ),
    RuleDefinition(
        "risk.sms.hourly",
        "短信设备每小时告警",
        "风控免告警",
        RISK_SMS_HOURLY_THRESHOLD_KEY,
        5,
        "北京时间自然小时",
        (SUBJECT_DEVICE,),
        max_value=100_000,
        allow_unlimited=False,
        alert_only=True,
    ),
    RuleDefinition(
        "risk.oneclick.daily",
        "一键登录设备每日告警",
        "风控免告警",
        RISK_ONECLICK_DAILY_THRESHOLD_KEY,
        20,
        "北京时间自然日",
        (SUBJECT_DEVICE,),
        max_value=100_000,
        allow_unlimited=False,
        alert_only=True,
    ),
    RuleDefinition(
        "risk.compare.daily",
        "比价账号每日告警",
        "风控免告警",
        RISK_COMPARE_DAILY_THRESHOLD_KEY,
        100,
        "北京时间自然日",
        (SUBJECT_PHONE,),
        max_value=100_000,
        allow_unlimited=False,
        alert_only=True,
    ),
)
RULE_MAP = {rule.code: rule for rule in RULES}
RULE_CODE_BY_CONFIG_KEY = {rule.config_key: rule.code for rule in RULES}
LIMIT_CONFIG_KEYS = tuple(RULE_CODE_BY_CONFIG_KEY)


@dataclass(frozen=True)
class EffectiveLimit:
    rule_code: str
    global_limit: int
    limit: int | None
    mode: str
    suppressed: bool
    override_id: int | None
    matched_subject_type: str | None
    matched_subject_value: str | None
    reset_at: datetime | None
    bucket_version: str

    @property
    def unlimited(self) -> bool:
        return self.limit is None


def normalize_subject(subject_type: str, value: str) -> str:
    value = (value or "").strip()
    if subject_type == SUBJECT_PHONE:
        value = "".join(ch for ch in value if ch.isdigit())
    if subject_type not in SUBJECT_TYPES:
        raise ValueError(f"unsupported subject type: {subject_type}")
    if not value:
        raise ValueError("subject value is empty")
    return value[:128]


def validate_whitelist_subject(subject_type: str, value: str) -> str:
    """Normalize a whitelist subject and reject unsafe pseudo-devices.

    Old clients without a device ID are grouped by public IP for rate limiting.
    That fallback remains valid at runtime, but it is not a stable, unique
    device identity and must never be persisted as a device whitelist target.
    """

    normalized = normalize_subject(subject_type, value)
    if (
        subject_type == SUBJECT_DEVICE
        and normalized.startswith(LEGACY_IP_DEVICE_PREFIX)
    ):
        raise ValueError(
            "旧客户端未上报真实设备 ID，不能加入设备白名单，请升级客户端后重试"
        )
    return normalized


def get_rule(rule_code: str) -> RuleDefinition:
    try:
        return RULE_MAP[rule_code]
    except KeyError as exc:
        raise ValueError(f"unknown rule: {rule_code}") from exc


def device_source_scope(rule_code: str) -> str:
    """返回设备候选数据所属命名空间，防止一个设备 ID 跨来源误套规则。"""
    rule = get_rule(rule_code)
    if SUBJECT_DEVICE not in rule.subject_types:
        raise ValueError("当前限制项不支持设备白名单")
    return "comparison" if rule_code == "compare.start.daily" else "auth"


def default_global_limits() -> dict[str, int]:
    """Return the complete 16-rule default snapshot keyed by rule code."""
    return {rule.code: rule.default_limit for rule in RULES}


def _normalise_global_limits(value: object) -> dict[str, int]:
    """Merge a stored JSON object with safe code defaults.

    The migration and every admin write persist all rules.  Defaults are still
    merged here so a manually damaged/older partial JSON cannot take the
    service down after deployment.
    """

    values = default_global_limits()
    if not isinstance(value, dict):
        return values
    for rule_code, raw in value.items():
        rule = RULE_MAP.get(str(rule_code))
        if rule is None or isinstance(raw, bool):
            continue
        try:
            parsed = int(raw)
        except (TypeError, ValueError):
            continue
        if rule.min_value <= parsed <= rule.max_value:
            values[rule.code] = parsed
    return values


def _legacy_global_limit(db: Session, rule: RuleDefinition) -> tuple[int, str]:
    """Read the pre-bundle representation while upgrading old/test databases."""

    row = db.get(AppConfig, rule.config_key)
    if row is not None:
        return int(row.value), "legacy-key"

    if rule.legacy_config_key:
        legacy = db.get(AppConfig, rule.legacy_config_key)
        if legacy is not None:
            value = legacy.value
            if rule.legacy_json_field:
                value = value.get(rule.legacy_json_field) if isinstance(value, dict) else None
            if value is not None:
                return int(value), "legacy"

    try:
        return int(app_config.get_value(db, rule.config_key)), "default"
    except KeyError:
        return rule.default_limit, "default"


def get_global_limits(db: Session) -> dict[str, int]:
    """Read the complete global-limit JSON, with a pre-migration fallback."""

    row = db.get(AppConfig, LIMIT_POLICY_GLOBAL_KEY)
    if row is not None:
        return _normalise_global_limits(row.value)
    return {rule.code: _legacy_global_limit(db, rule)[0] for rule in RULES}


def set_global_limits(
    db: Session,
    updates: dict[str, int],
    *,
    admin_id: int,
    commit: bool = True,
) -> dict[str, int]:
    """Atomically update selected rules inside the single complete JSON row."""

    parsed_updates: dict[str, int] = {}
    for rule_code, raw_value in updates.items():
        rule = get_rule(rule_code)
        value = int(raw_value)
        if not rule.min_value <= value <= rule.max_value:
            raise ValueError(
                f"limit for {rule_code} must be between "
                f"{rule.min_value} and {rule.max_value}"
            )
        parsed_updates[rule.code] = value

    row = db.scalar(
        select(AppConfig)
        .where(AppConfig.key == LIMIT_POLICY_GLOBAL_KEY)
        .with_for_update()
    )
    values = (
        _normalise_global_limits(row.value)
        if row is not None
        else {rule.code: _legacy_global_limit(db, rule)[0] for rule in RULES}
    )
    values.update(parsed_updates)

    if row is None:
        row = AppConfig(
            key=LIMIT_POLICY_GLOBAL_KEY,
            value=values,
            updated_by_admin_id=admin_id,
        )
        db.add(row)
    else:
        row.value = dict(values)
        row.updated_by_admin_id = admin_id

    # Once the bundle exists, stale sparse rows must not become a second source
    # of truth.  The data migration performs the same cleanup for production.
    db.execute(delete(AppConfig).where(AppConfig.key.in_(LIMIT_CONFIG_KEYS)))
    if commit:
        db.commit()
        db.refresh(row)
    else:
        db.flush()
    return values


def _global_limit(db: Session, rule: RuleDefinition) -> tuple[int, str]:
    row = db.get(AppConfig, LIMIT_POLICY_GLOBAL_KEY)
    if row is not None:
        return _normalise_global_limits(row.value)[rule.code], "configured"
    return _legacy_global_limit(db, rule)


def _aware(value: datetime | None) -> datetime | None:
    if value is None:
        return None
    return value.replace(tzinfo=UTC) if value.tzinfo is None else value


def _matching_overrides(
    db: Session,
    rule: RuleDefinition,
    subjects: dict[str, str | None],
    now: datetime,
) -> list[LimitPolicyOverride]:
    pairs: list[tuple[str, str]] = []
    for subject_type in rule.subject_types:
        raw = subjects.get(subject_type)
        if raw:
            pairs.append((subject_type, normalize_subject(subject_type, raw)))
    if not pairs:
        return []
    clauses = [
        (
            (LimitPolicyOverride.subject_type == subject_type)
            & (LimitPolicyOverride.subject_value == subject_value)
        )
        for subject_type, subject_value in pairs
    ]
    rows = list(
        db.execute(
            select(LimitPolicyOverride).where(
                LimitPolicyOverride.rule_code == rule.code,
                LimitPolicyOverride.enabled.is_(True),
                or_(*clauses),
            )
        ).scalars()
    )
    active = [
        row
        for row in rows
        if (_aware(row.starts_at) is None or _aware(row.starts_at) <= now)
        and (_aware(row.expires_at) is None or _aware(row.expires_at) > now)
    ]
    return sorted(active, key=lambda row: SUBJECT_PRECEDENCE[row.subject_type])


def resolve(
    db: Session,
    rule_code: str,
    *,
    phone: str | None = None,
    device: str | None = None,
    now: datetime | None = None,
) -> EffectiveLimit:
    """Resolve global config plus the most specific active override.

    Device overrides win over phone overrides when both match.
    """

    rule = get_rule(rule_code)
    now = _aware(now) or datetime.now(UTC)
    global_limit, global_version = _global_limit(db, rule)
    matches = _matching_overrides(
        db, rule, {SUBJECT_PHONE: phone, SUBJECT_DEVICE: device}, now
    )
    row = matches[0] if matches else None
    if row is None:
        return EffectiveLimit(
            rule.code,
            global_limit,
            global_limit,
            MODE_INHERIT,
            False,
            None,
            None,
            None,
            None,
            global_version,
        )

    limit: int | None = global_limit
    suppressed = False
    if row.mode == MODE_OVERRIDE:
        limit = int(row.limit_value) if row.limit_value is not None else global_limit
    elif row.mode == MODE_UNLIMITED:
        limit = None
    elif row.mode == MODE_SUPPRESS_ALERT:
        suppressed = True
    # 编辑限制值/备注不应隐式清空计数；只有显式“重置状态”才切换桶。
    reset_version = _aware(row.reset_at)
    version = (
        f"{global_version}:o:{row.id}:"
        f"{reset_version.isoformat() if reset_version else '0'}"
    )
    return EffectiveLimit(
        rule.code,
        global_limit,
        limit,
        row.mode,
        suppressed,
        row.id,
        row.subject_type,
        row.subject_value,
        _aware(row.reset_at),
        version,
    )


def resolve_for_user(
    db: Session,
    rule_code: str,
    user_id: int,
    *,
    device: str | None = None,
) -> EffectiveLimit:
    user = db.get(User, user_id)
    return resolve(
        db,
        rule_code,
        phone=user.phone if user is not None else None,
        device=device,
    )


def rule_catalog(db: Session) -> list[dict]:
    out: list[dict] = []
    values = get_global_limits(db)
    for rule in RULES:
        out.append(
            {
                "code": rule.code,
                "label": rule.label,
                "group": rule.group,
                "global_limit": values[rule.code],
                "default_limit": rule.default_limit,
                "window_label": rule.window_label,
                "subject_types": list(rule.subject_types),
                "allowed_modes": list(rule.allowed_modes),
                "min_value": rule.min_value,
                "max_value": rule.max_value,
                "supports_reset": rule.supports_reset,
                "alert_only": rule.alert_only,
            }
        )
    return out


def validate_override(
    rule: RuleDefinition,
    *,
    subject_type: str,
    mode: str,
    limit_value: int | None,
    starts_at: datetime | None,
    expires_at: datetime | None,
) -> None:
    if subject_type not in rule.subject_types:
        raise ValueError("该规则不支持此主体类型")
    if mode not in rule.allowed_modes:
        raise ValueError("该规则不支持此策略模式")
    if limit_value is not None:
        raise ValueError("白名单不支持覆盖指定值")
    if mode in {MODE_UNLIMITED, MODE_SUPPRESS_ALERT} and expires_at is None:
        raise ValueError("临时白名单必须设置失效时间")
    if starts_at and expires_at and _aware(expires_at) <= _aware(starts_at):
        raise ValueError("失效时间必须晚于生效时间")
    if (
        mode in {MODE_UNLIMITED, MODE_SUPPRESS_ALERT}
        and expires_at is not None
        and _aware(expires_at) <= datetime.now(UTC)
    ):
        raise ValueError("临时白名单的失效时间必须晚于当前时间")
    if mode == MODE_SUPPRESS_ALERT and not rule.alert_only:
        raise ValueError("免告警只支持风控监控的三项规则")


def active_overrides(
    db: Session,
    *,
    subject_type: str | None = None,
    keyword: str | None = None,
) -> Iterable[LimitPolicyOverride]:
    stmt = select(LimitPolicyOverride)
    if subject_type:
        stmt = stmt.where(LimitPolicyOverride.subject_type == subject_type)
    if keyword:
        stmt = stmt.where(LimitPolicyOverride.subject_value.ilike(f"%{keyword.strip()}%"))
    return db.execute(
        stmt.order_by(LimitPolicyOverride.updated_at.desc(), LimitPolicyOverride.id.desc())
    ).scalars()