|
|
|
|
@ -0,0 +1,132 @@
|
|
|
|
|
"""
|
|
|
|
|
验证码服务（Redis 存储 + Flask-Mail 发送）
|
|
|
|
|
|
|
|
|
|
提供：
|
|
|
|
|
- `send_verification_code(email, purpose='register', length=6, expire_seconds=None)`
|
|
|
|
|
- `verify_code(email, code, purpose='register')`
|
|
|
|
|
|
|
|
|
|
依赖：
|
|
|
|
|
- `redis`（通过 `REDIS_URL` 配置，默认为 redis://localhost:6379/0）
|
|
|
|
|
- Flask-Mail 已在应用中初始化（通过 `current_app.extensions['mail']` 获取）
|
|
|
|
|
"""
|
|
|
|
|
import random
|
|
|
|
|
import string
|
|
|
|
|
import logging
|
|
|
|
|
from typing import Optional
|
|
|
|
|
|
|
|
|
|
import redis
|
|
|
|
|
from flask import current_app
|
|
|
|
|
from flask_mail import Message
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
pool = redis.ConnectionPool().from_url('redis://localhost:6379/0', decode_responses=True)
|
|
|
|
|
|
|
|
|
|
def _get_redis_client() -> redis.Redis:
|
|
|
|
|
"""根据 REDIS_URL 创建 redis 客户端"""
|
|
|
|
|
redis_url = current_app.config.get('REDIS_URL', 'redis://localhost:6379/0')
|
|
|
|
|
|
|
|
|
|
# 不再使用全局的 pool，而是每次根据 URL 获取连接（redis-py 内部自己会管理连接池）
|
|
|
|
|
# 或者如果你想复用连接池，应该在 app 启动时初始化 pool，而不是在全局写死 localhost
|
|
|
|
|
return redis.Redis.from_url(redis_url, decode_responses=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _generate_code(length: int = 6) -> str:
|
|
|
|
|
"""生成指定长度的数字验证码（默认 6 位）。"""
|
|
|
|
|
# 只产生数字字符串更常用于验证码
|
|
|
|
|
return ''.join(random.choices(string.digits, k=length))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def send_verification_code(email: str,
|
|
|
|
|
purpose: str = 'register',
|
|
|
|
|
length: int = 6,
|
|
|
|
|
expire_seconds: Optional[int] = None) -> bool:
|
|
|
|
|
"""生成验证码，保存到 Redis，并使用 Flask-Mail 发送给 `email`。
|
|
|
|
|
|
|
|
|
|
返回 True 表示发送成功，False 表示失败（或抛出异常时捕获后返回 False）。
|
|
|
|
|
"""
|
|
|
|
|
# 读取过期时间（秒），优先使用传入值，其次使用 app config
|
|
|
|
|
if expire_seconds is None:
|
|
|
|
|
expire_seconds = current_app.config.get('VERIFICATION_CODE_EXPIRES', 300)
|
|
|
|
|
|
|
|
|
|
code = _generate_code(length)
|
|
|
|
|
key = f"verify:{purpose}:{email}"
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
r = _get_redis_client()
|
|
|
|
|
# 使用字符串保存验证码，并设置过期时间
|
|
|
|
|
r.set(key, code, ex=expire_seconds)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logger.exception("保存验证码到 Redis 失败")
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
# 尝试发送邮件
|
|
|
|
|
try:
|
|
|
|
|
mail = current_app.extensions.get('mail')
|
|
|
|
|
subject = "您的验证码" # current_app.config.get('VERIFICATION_EMAIL_SUBJECT', '您的验证码')
|
|
|
|
|
sender = '1798231811@qq.com' # current_app.config.get('MAIL_DEFAULT_SENDER') or current_app.config.get('MAIL_USERNAME')
|
|
|
|
|
body = f'您的验证码为：{code}，有效期 {expire_seconds} 秒。' # current_app.config.get('VERIFICATION_EMAIL_TEMPLATE',
|
|
|
|
|
# f'您的验证码为：{code}，有效期 {expire_seconds} 秒。')
|
|
|
|
|
|
|
|
|
|
# 优先使用简单文本邮件，项目中可按需替换为 HTML 模板
|
|
|
|
|
msg = Message(subject=subject, recipients=[email], body=body, sender=sender)
|
|
|
|
|
|
|
|
|
|
if mail is None:
|
|
|
|
|
# 如果 Flask-Mail 未初始化，记录日志并返回 False
|
|
|
|
|
logger.error('Flask-Mail 未初始化，无法发送邮件')
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
mail.send(msg)
|
|
|
|
|
logger.info('已发送验证码到 %s (purpose=%s)', email, purpose)
|
|
|
|
|
return True
|
|
|
|
|
except Exception:
|
|
|
|
|
logger.exception('发送验证码邮件失败')
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def verify_code(email: str, code: str, purpose: str = 'register') -> bool:
|
|
|
|
|
"""校验验证码是否正确。成功可配置是否从 Redis 删除该 key。
|
|
|
|
|
|
|
|
|
|
返回 True 表示校验通过；False 表示失败或异常。
|
|
|
|
|
"""
|
|
|
|
|
key = f"verify:{purpose}:{email}"
|
|
|
|
|
try:
|
|
|
|
|
r = _get_redis_client()
|
|
|
|
|
stored = r.get(key)
|
|
|
|
|
if stored is None:
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
matched = (str(stored) == str(code))
|
|
|
|
|
if matched :
|
|
|
|
|
try:
|
|
|
|
|
r.delete(key)
|
|
|
|
|
except Exception:
|
|
|
|
|
logger.warning('校验成功，但删除 Redis key 失败: %s', key)
|
|
|
|
|
return matched
|
|
|
|
|
except Exception:
|
|
|
|
|
logger.exception('校验验证码时发生异常')
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def clear_verification_code(email: str, purpose: str = 'register') -> bool:
|
|
|
|
|
"""显式删除指定 email 的验证码（例如用于管理员撤销）。"""
|
|
|
|
|
key = f"verify:{purpose}:{email}"
|
|
|
|
|
try:
|
|
|
|
|
r = _get_redis_client()
|
|
|
|
|
return r.delete(key) == 1
|
|
|
|
|
except Exception:
|
|
|
|
|
logger.exception('删除验证码失败')
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
# 简单测试发送和验证功能
|
|
|
|
|
test_email = "3310207578@qq.com"
|
|
|
|
|
if send_verification_code(test_email, expire_seconds=600):
|
|
|
|
|
print("验证码发送成功")
|
|
|
|
|
code = input("请输入收到的验证码: ")
|
|
|
|
|
if verify_code(test_email, code):
|
|
|
|
|
print("验证码验证成功")
|
|
|
|
|
else:
|
|
|
|
|
print("验证码验证失败")
|
|
|
|
|
else:
|
|
|
|
|
print("验证码发送失败")
|
