火山引擎SMS API|celery容灾优化|生成模型引擎积分列表API

2026-06-04 16:28:59 +08:00
parent 450e509b1a
commit 8d43d5871c
25 changed files with 1936 additions and 302 deletions
@@ -1,79 +1,186 @@
+import asyncio
+import json
 import logging
 import random
 import time
-
-import httpx
+from datetime import datetime
+from typing import Any

 from app.config import settings
 from app.utils.redis import get_redis

 logger = logging.getLogger("videogen")

-# In-memory fallback for verification codes
+# In-memory fallback for local development when Redis is disabled.
 _sms_code_store: dict[str, tuple[str, float]] = {}
+_sms_send_interval_store: dict[str, float] = {}
+_sms_daily_count_store: dict[str, tuple[str, int]] = {}


-def _generate_code(length: int = 6) -> str:
-    return "".join(random.choices("0123456789", k=length))
+def _generate_code(length: int | None = None) -> str:
+    code_length = int(length or settings.SMS_CODE_LENGTH or 4)
+    code_length = max(4, min(code_length, 8))
+    return "".join(random.choices("0123456789", k=code_length))


-async def send_sms(phone: str, code: str) -> bool:
-    """Send SMS verification code. Supports mock mode and real HTTP gateway."""
-    if settings.SMS_MOCK or not settings.SMS_API_URL:
-        logger.info(f"[SMS MOCK] To={phone}, Code={code}")
-        return True
+def _normalize_scene(scene: str | None) -> str:
+    value = (scene or "").strip().lower()
+    if value not in {"register", "login", "set_password"}:
+        value = "login"
+    return value
+
+
+def _code_key(phone: str, scene: str | None) -> str:
+    return f"sms_code:{_normalize_scene(scene)}:{phone}"
+
+
+def _interval_key(phone: str, scene: str | None) -> str:
+    return f"sms_interval:{_normalize_scene(scene)}:{phone}"
+
+
+def _daily_key(phone: str, scene: str | None) -> str:
+    today = datetime.now().strftime("%Y%m%d")
+    return f"sms_daily:{_normalize_scene(scene)}:{phone}:{today}"
+
+
+def _today() -> str:
+    return datetime.now().strftime("%Y%m%d")
+
+
+async def _check_send_limit(phone: str, scene: str) -> None:
+    """发送频控：同场景同手机号间隔限制 + 每日次数限制。"""
+    redis = get_redis()
+    interval_seconds = int(settings.SMS_SEND_INTERVAL_SECONDS or 60)
+    daily_limit = int(settings.SMS_DAILY_LIMIT or 20)
+
+    if redis:
+        interval_key = _interval_key(phone, scene)
+        if await redis.get(interval_key):
+            raise ValueError(f"短信发送过于频繁，请{interval_seconds}秒后再试")
+
+        daily_key = _daily_key(phone, scene)
+        count = await redis.incr(daily_key)
+        if count == 1:
+            await redis.expire(daily_key, 24 * 60 * 60)
+        if count > daily_limit:
+            raise ValueError("今日短信发送次数已达上限，请明天再试")
+
+        await redis.setex(interval_key, interval_seconds, "1")
+        return
+
+    now_ts = time.time()
+    interval_key = _interval_key(phone, scene)
+    last_send_at = _sms_send_interval_store.get(interval_key)
+    if last_send_at and now_ts - last_send_at < interval_seconds:
+        raise ValueError(f"短信发送过于频繁，请{interval_seconds}秒后再试")
+    _sms_send_interval_store[interval_key] = now_ts
+
+    daily_key = f"{_normalize_scene(scene)}:{phone}"
+    current_day = _today()
+    stored_day, count = _sms_daily_count_store.get(daily_key, (current_day, 0))
+    if stored_day != current_day:
+        stored_day, count = current_day, 0
+    count += 1
+    _sms_daily_count_store[daily_key] = (stored_day, count)
+    if count > daily_limit:
+        raise ValueError("今日短信发送次数已达上限，请明天再试")
+
+
+def _send_volc_sms_sync(phone: str, code: str) -> dict[str, Any]:
+    from volcengine.sms.SmsService import SmsService
+
+    if not settings.VOLC_SMS_ACCESS_KEY_ID or not settings.VOLC_SMS_SECRET_ACCESS_KEY:
+        raise RuntimeError("火山短信 AK/SK 未配置")
+    if not settings.VOLC_SMS_ACCOUNT:
+        raise RuntimeError("火山短信消息组ID VOLC_SMS_ACCOUNT 未配置")
+    if not settings.VOLC_SMS_TEMPLATE_ID:
+        raise RuntimeError("火山短信模板ID VOLC_SMS_TEMPLATE_ID 未配置")
+    if not settings.VOLC_SMS_SIGN:
+        raise RuntimeError("火山短信签名 VOLC_SMS_SIGN 未配置")
+
+    sms_service = SmsService()
+    sms_service.set_ak(settings.VOLC_SMS_ACCESS_KEY_ID)
+    sms_service.set_sk(settings.VOLC_SMS_SECRET_ACCESS_KEY)
+
+    body = {
+        "SmsAccount": settings.VOLC_SMS_ACCOUNT,
+        "Sign": settings.VOLC_SMS_SIGN,
+        "TemplateID": settings.VOLC_SMS_TEMPLATE_ID,
+        "TemplateParam": json.dumps({"xxxx": str(code)}, ensure_ascii=False, separators=(",", ":")),
+        "Tag": f"{phone}:{int(time.time())}",
+        "PhoneNumbers": phone,
+    }
+    raw_resp = sms_service.send_sms(json.dumps(body, ensure_ascii=False, separators=(",", ":")))
+
+    if isinstance(raw_resp, str):
+        try:
+            resp: dict[str, Any] = json.loads(raw_resp)
+        except json.JSONDecodeError:
+            resp = {"raw": raw_resp}
+    elif isinstance(raw_resp, dict):
+        resp = raw_resp
+    else:
+        resp = {"raw": raw_resp}
+
+    error = (resp.get("ResponseMetadata") or {}).get("Error") if isinstance(resp, dict) else None
+    if error:
+        raise RuntimeError(f"火山短信发送失败：{error.get('Code')} {error.get('Message')}")
+
+    return resp
+
+
+async def send_sms(phone: str, code: str, scene: str = "login") -> bool:
+    """发送短信验证码。
+
+    SMS_MOCK=true 时只写日志，方便本地调试；否则使用火山引擎短信 SDK。
+    火山 SDK 为同步调用，这里放到线程中执行，避免阻塞 FastAPI event loop。
+    """
+    scene = _normalize_scene(scene)

    try:
-        payload = {
-            "phone": phone,
-            "code": code,
-            "sign_name": settings.SMS_SIGN_NAME,
-            "template_code": settings.SMS_TEMPLATE_CODE,
-        }
-        headers = {
-            "Authorization": f"Bearer {settings.SMS_API_KEY}",
-            "Content-Type": "application/json",
-        }
-        async with httpx.AsyncClient(timeout=10) as client:
-            resp = await client.post(settings.SMS_API_URL, json=payload, headers=headers)
-            resp.raise_for_status()
-            return True
+        logger.info("[SMS] scene=%s, To=%s, Code=%s", scene, phone, code)
+        await asyncio.to_thread(_send_volc_sms_sync, phone, code)
+        return True
    except Exception:
-        logger.exception(f"SMS send failed for {phone}")
+        logger.exception("SMS send failed. scene=%s phone=%s", scene, phone)
        return False


-async def store_sms_code(phone: str, code: str, ttl: int = 300) -> None:
-    """Store SMS verification code with TTL (default 5 minutes)."""
+async def store_sms_code(phone: str, code: str, scene: str = "login", ttl: int | None = None) -> None:
+    ttl_seconds = int(ttl or settings.SMS_CODE_TTL_SECONDS or 300)
+    key = _code_key(phone, scene)
    redis = get_redis()
    if redis:
-        await redis.setex(f"sms_code:{phone}", ttl, code)
+        await redis.setex(key, ttl_seconds, code)
    else:
-        _sms_code_store[phone] = (code, time.time() + ttl)
+        _sms_code_store[key] = (code, time.time() + ttl_seconds)


-async def generate_and_send_sms(phone: str) -> bool:
-    """Generate a code, store it, and send it via SMS."""
+async def generate_and_send_sms(phone: str, scene: str = "login") -> bool:
+    scene = _normalize_scene(scene)
+    await _check_send_limit(phone, scene)
    code = _generate_code()
-    ok = await send_sms(phone, code)
+    ok = await send_sms(phone, code, scene)
    if ok:
-        await store_sms_code(phone, code)
+        await store_sms_code(phone, code, scene)
    return ok


-async def verify_sms_code(phone: str, code: str) -> bool:
-    """Verify an SMS verification code."""
+async def verify_sms_code(phone: str, code: str, scene: str = "login") -> bool:
+    key = _code_key(phone, scene)
    redis = get_redis()
    if redis:
-        stored = await redis.get(f"sms_code:{phone}")
-        if stored and stored == code:
-            await redis.delete(f"sms_code:{phone}")
+        stored = await redis.get(key)
+        if isinstance(stored, bytes):
+            stored = stored.decode()
+        if stored and str(stored) == str(code):
+            await redis.delete(key)
            return True
        return False
-    else:
-        entry = _sms_code_store.pop(phone, None)
-        if entry:
-            stored_code, expires = entry
-            if time.time() < expires and stored_code == code:
-                return True
-        return False
+
+    entry = _sms_code_store.pop(key, None)
+    if entry:
+        stored_code, expires = entry
+        if time.time() < expires and stored_code == code:
+            return True
+    return False