verify_code_engine.py

# verify_code_engine.py
# ==========================================
# 📧 验证码发送引擎
# ==========================================
# 作用：提供验证码内存缓存与多渠道发送能力（邮件/短信）
# 关联文件：
#   - router_users_auth.py (调用此模块发送验证码)
#   - router_wallet.py (提现时也会调用验证码缓存)
# ==========================================

import os
import json
import time
import urllib.request
import urllib.parse
import base64

from fastapi import HTTPException

# ==========================================
# 验证码内存缓存 (全局共享字典)
# ==========================================
# 作用：存储所有待验证的验证码，格式为 {key: {code, expires_at}}
# 关联：router_users_auth.py 和 router_wallet.py 都会读写此字典
# 注意：服务重启后缓存会清空，生产环境建议改用 Redis
VERIFY_CODES = {}

# 🔒 发送频率限制配置
SEND_COOLDOWN = {}  # {contact: last_send_timestamp}
COOLDOWN_SECONDS = 60  # 同一联系方式60秒内最多发送1条验证码

# 🚀 P1性能优化：记录上次清理时间，避免频繁清理
_last_cleanup_time = 0
_CLEANUP_INTERVAL = 300  # 5分钟清理一次


def check_send_cooldown(contact: str):
    """
    🔒 检查发送频率限制
    
    作用：防止同一联系方式频繁发送验证码
    参数：
        - contact: 联系方式（邮箱或手机号）
    抛出：
        - HTTPException: 如果在冷却期内，返回429状态码
    """
    now = time.time()
    last_send = SEND_COOLDOWN.get(contact, 0)
    if now - last_send < COOLDOWN_SECONDS:
        remaining = int(COOLDOWN_SECONDS - (now - last_send))
        raise HTTPException(status_code=429, detail=f"请等待{remaining}秒后再发送验证码")
    SEND_COOLDOWN[contact] = now


def cleanup_expired_codes():
    """
    🚀 P1性能优化：清理过期验证码，防止内存泄漏
    🔒 同时清理过期的发送频率限制记录
    """
    global _last_cleanup_time
    now = time.time()
    
    # 限制清理频率，避免每次请求都清理
    if now - _last_cleanup_time < _CLEANUP_INTERVAL:
        return
    
    _last_cleanup_time = now
    
    # 收集过期的验证码键
    expired_keys = [
        key for key, data in VERIFY_CODES.items()
        if now > data.get("expires_at", data.get("expires", 0))
    ]
    
    # 删除过期验证码
    for key in expired_keys:
        del VERIFY_CODES[key]
    
    # 🔒 清理过期的发送频率限制记录，防止内存泄漏
    expired_cooldown_contacts = [
        contact for contact, last_send in SEND_COOLDOWN.items()
        if now - last_send >= COOLDOWN_SECONDS
    ]
    for contact in expired_cooldown_contacts:
        del SEND_COOLDOWN[contact]
    
    total_cleaned = len(expired_keys) + len(expired_cooldown_contacts)
    if total_cleaned > 0:
        print(f"🧹 已清理 {len(expired_keys)} 个过期验证码，{len(expired_cooldown_contacts)} 个过期冷却记录，当前缓存: {len(VERIFY_CODES)} 条验证码，{len(SEND_COOLDOWN)} 条冷却记录")


def send_email_code(to_email: str, code: str, action: str):
    """
    📧 邮件验证码发送函数
    
    作用：通过 Make.com Webhook 触发邮件发送（无需 SMTP 服务器）
    参数：
        - to_email: 目标邮箱地址
        - code: 6位数字验证码
        - action: 动作类型 ("register" 注册 / "reset" 重置密码 / "withdraw" 提现)
    关联：
        - 环境变量 MAKE_WEBHOOK_URL (在 HF Space Settings 中配置)
        - router_users_auth.py 的 send_verify_code() 异步调用此函数
    """
    # 从环境变量读取 Make.com Webhook 地址
    webhook_url = os.environ.get("MAKE_WEBHOOK_URL")

    if not webhook_url:
        print("警告: 未配置 MAKE_WEBHOOK_URL，跳过邮件发送")
        return

    # 根据动作类型生成不同的邮件标题
    if action == "register":
        action_str = "注册账号"
    elif action == "withdraw":
        action_str = "提现操作"
    else:
        action_str = "修改/找回密码"
    subject = f"ComfyUI 社区 - {action_str}验证码"
    
    # 构建 HTML 格式的邮件正文（美化样式）
    html_content = f"""
    <div style="background:#f9f9f9; padding:20px; font-family:sans-serif;">
        <div style="background:#fff; padding:20px; border-radius:8px; max-width:500px; margin:0 auto; box-shadow:0 2px 10px rgba(0,0,0,0.05);">
            <h2 style="color:#4CAF50; margin-top:0;">ComfyUI 社区精选</h2>
            <p>您好，</p>
            <p>您正在请求<strong>{action_str}</strong>，您的验证码是：</p>
            <div style="font-size:24px; font-weight:bold; color:#2196F3; background:#e3f2fd; padding:15px; text-align:center; border-radius:6px; letter-spacing: 5px;">{code}</div>
            <p style="color:#888; font-size:12px; margin-top:20px;">该验证码在 5 分钟内有效。如非本人操作，请忽略此邮件。</p>
        </div>
    </div>
    """

    # 组装 Webhook 请求载荷
    data = {
        "to": to_email,
        "subject": subject,
        "html": html_content
    }

    # 发送 HTTP POST 请求触发 Webhook
    req = urllib.request.Request(
        webhook_url, 
        data=json.dumps(data).encode('utf-8'), 
        headers={'Content-Type': 'application/json'}
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as response:
            print(f"✅ 成功触发 Webhook 发送验证码 {code} 至 {to_email}")
    except Exception as e:
        print(f"❌ Webhook 触发失败: {e}")


def send_sms_code(phone: str, code: str, action: str):
    """
    📱 短信验证码发送函数（双引擎：Twilio + 阿里云）
    
    作用：支持海外用户(Twilio)和国内用户(阿里云)的短信发送
    参数：
        - phone: 手机号（需带国际区号，如 +86）
        - code: 6位数字验证码
        - action: 动作类型 ("register" 注册 / "reset" 重置密码 / "withdraw" 提现)
    关联：
        - 环境变量 TWILIO_SID, TWILIO_TOKEN, TWILIO_FROM (Twilio配置)
        - 环境变量 ALIYUN_AK, ALIYUN_SK, ALIYUN_SIGN_NAME (阿里云配置)
        - router_users_auth.py 的 send_verify_code() 异步调用此函数
    """
    if action == "register":
        action_str = "注册账号"
    elif action == "withdraw":
        action_str = "提现操作"
    else:
        action_str = "修改/找回密码"

    # ==========================================
    # 引擎 A：Twilio (海外优先，无需SDK，纯HTTP接口)
    # ==========================================
    # 作用：为海外用户或测试环境提供短信能力
    # 关联：需在 HF Space Settings 配置 TWILIO_SID, TWILIO_TOKEN, TWILIO_FROM
    twilio_sid = os.environ.get("TWILIO_SID")
    if twilio_sid:
        token = os.environ.get("TWILIO_TOKEN")
        from_phone = os.environ.get("TWILIO_FROM")
        
        # 构建短信内容
        body = f"【ComfyUI社区】您正在请求{action_str}，验证码是：{code}，5分钟内有效。"
        url = f"https://api.twilio.com/2010-04-01/Accounts/{twilio_sid}/Messages.json"
        
        # Twilio 使用 Basic Auth 认证
        auth = base64.b64encode(f"{twilio_sid}:{token}".encode('utf-8')).decode('utf-8')
        data = urllib.parse.urlencode({'To': phone, 'From': from_phone, 'Body': body}).encode('utf-8')
        
        req = urllib.request.Request(url, data=data)
        req.add_header("Authorization", f"Basic {auth}")
        try:
            with urllib.request.urlopen(req, timeout=10) as response:
                print(f"✅ Twilio 短信已成功下发至 {phone}")
        except Exception as e:
            print(f"❌ Twilio 发送失败: {e}")
        return

    # ==========================================
    # 引擎 B：阿里云 (国内首选，到达率最高)
    # ==========================================
    # 作用：为国内用户提供短信能力，需要已备案的签名和模板
    # 关联：需在 HF Space Settings 配置 ALIYUN_AK, ALIYUN_SK 等
    aliyun_ak = os.environ.get("ALIYUN_AK")
    if aliyun_ak:
        # 阿里云SDK需要单独安装，这里做动态导入
        try:
            from alibabacloud_dysmsapi20170525.client import Client as DysmsapiClient
            from alibabacloud_tea_openapi import models as open_api_models
            from alibabacloud_dysmsapi20170525 import models as dysmsapi_models
        except ImportError:
            print("❌ 缺少阿里云 SDK，请在 requirements.txt 中添加 alibabacloud_dysmsapi20170525")
            return

        sk = os.environ.get("ALIYUN_SK")
        sign_name = os.environ.get("ALIYUN_SIGN_NAME")  # 短信签名，如 "阿里云"
        # 根据动作类型选择不同的模板
        tpl_code = os.environ.get("ALIYUN_TPL_REGISTER") if action == "register" else os.environ.get("ALIYUN_TPL_RESET")

        # 初始化阿里云客户端
        config = open_api_models.Config(access_key_id=aliyun_ak, access_key_secret=sk)
        config.endpoint = 'dysmsapi.aliyuncs.com'
        client = DysmsapiClient(config)

        # 构建发送请求
        send_req = dysmsapi_models.SendSmsRequest(
            phone_numbers=phone,
            sign_name=sign_name,
            template_code=tpl_code,
            template_param=json.dumps({"code": code})  # 模板参数
        )
        try:
            response = client.send_sms(send_req)
            if response.body.code == "OK":
                print(f"✅ 阿里云短信已成功下发至 {phone}")
            else:
                print(f"❌ 阿里云下发失败: {response.body.message}")
        except Exception as e:
            print(f"❌ 阿里云请求异常: {e}")
        return

    # ==========================================
    # 降级模式：控制台打印模拟
    # ==========================================
    # 作用：当没有配置任何短信服务时，仅在控制台打印验证码（仅限开发测试）
    print(f"⚠️ 未配置短信秘钥，模拟下发 -> 手机号: {phone}, 验证码: {code}")