forked from xiaohei/taiji-AI-PAD
chenchenandClaude Opus 4.7
eb17ed84f8
A failed POST /api/auth/register returned the SQLAlchemy IntegrityError verbatim
to the caller, which included the full INSERT INTO tenant_model_keys statement
along with every bound parameter — ~50 plaintext LiteLLM API keys per failed
attempt. Same pattern was reproduced in 3 channel.py endpoints that wrap
LiteLLM key INSERTs.
Changes:
- channel.py: assign_resources_to_tenant / assign_model_to_tenant /
update_tenant_model_quota — log full exc_info, return a typed
{code, message} error instead of f"...{str(e)}". 6 leakage points sealed.
- email_verification.py: add peek_verification_code() — checks a code
without burning it. Lets the register handler verify *before* the
multi-step transaction so a downstream failure doesn't waste the user's
one-shot code.
- scripts/cleanup_orphan_litellm_keys.py: one-shot orphan key reaper.
Scans LiteLLM /key/list by metadata.tenant_id (plus a manual list of
the 8 publicly-leaked sk- prefixes from the original incident).
Used to nuke 16 orphan keys for tenant fab9dc27-… on 2026-05-12.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
427 lines
14 KiB
Python
427 lines
14 KiB
Python
"""
|
||
邮箱验证码功能
|
||
"""
|
||
|
||
import os
|
||
import random
|
||
import string
|
||
import smtplib
|
||
from email.mime.text import MIMEText
|
||
from email.mime.multipart import MIMEMultipart
|
||
from typing import Optional
|
||
import structlog
|
||
from app.state import get_state
|
||
|
||
logger = structlog.get_logger(__name__)
|
||
|
||
# 邮箱配置 - 从环境变量读取
|
||
SMTP_SERVER = os.getenv("SMTP_SERVER", "smtp.189.cn")
|
||
SMTP_PORT = int(os.getenv("SMTP_PORT", "465"))
|
||
SMTP_EMAIL = os.getenv("SMTP_EMAIL", "taijiagent@189.cn")
|
||
# 注意：生产环境必须通过环境变量设置SMTP_PASSWORD
|
||
# 如果未设置，邮件发送功能将不可用（测试模式下可通过日志或Redis获取验证码）
|
||
SMTP_PASSWORD = os.getenv("SMTP_PASSWORD")
|
||
# 是否使用SSL（端口465使用SSL，端口587使用STARTTLS）
|
||
SMTP_USE_SSL = os.getenv("SMTP_USE_SSL", "true").lower() == "true"
|
||
|
||
# 验证码配置
|
||
VERIFICATION_CODE_LENGTH = 6
|
||
VERIFICATION_CODE_EXPIRE_SECONDS = 600 # 10分钟
|
||
VERIFICATION_CODE_RATE_LIMIT_SECONDS = 60 # 发送频率限制：60秒内只能发送一次
|
||
|
||
|
||
def _send_email_sync(msg: MIMEMultipart) -> None:
|
||
"""同步发送邮件（在 executor 中运行）"""
|
||
if not SMTP_PASSWORD:
|
||
raise ValueError("SMTP_PASSWORD环境变量未设置，无法发送邮件")
|
||
|
||
# 根据端口选择连接方式：465用SSL，587用STARTTLS
|
||
if SMTP_PORT == 465 or SMTP_USE_SSL:
|
||
server = smtplib.SMTP_SSL(SMTP_SERVER, SMTP_PORT)
|
||
else:
|
||
server = smtplib.SMTP(SMTP_SERVER, SMTP_PORT)
|
||
server.starttls()
|
||
|
||
try:
|
||
server.login(SMTP_EMAIL, SMTP_PASSWORD)
|
||
server.send_message(msg)
|
||
logger.info("邮件发送成功", to=msg['To'])
|
||
except smtplib.SMTPAuthenticationError as e:
|
||
logger.error("SMTP认证失败", error=str(e), smtp_server=SMTP_SERVER, smtp_email=SMTP_EMAIL)
|
||
raise
|
||
except smtplib.SMTPException as e:
|
||
logger.error("SMTP发送失败", error=str(e))
|
||
raise
|
||
finally:
|
||
server.quit()
|
||
|
||
|
||
def generate_verification_code() -> str:
|
||
"""生成6位数字验证码"""
|
||
return ''.join(random.choices(string.digits, k=VERIFICATION_CODE_LENGTH))
|
||
|
||
|
||
async def send_verification_code(email: str, code: str, purpose: str = "register") -> bool:
|
||
"""
|
||
发送验证码邮件
|
||
|
||
Args:
|
||
email: 收件人邮箱
|
||
code: 验证码
|
||
purpose: 用途，"register" 表示注册，"reset_password" 表示重置密码
|
||
|
||
Returns:
|
||
是否发送成功
|
||
"""
|
||
if not SMTP_PASSWORD:
|
||
logger.error("SMTP密码未配置，无法发送邮件", email=email)
|
||
return False
|
||
|
||
try:
|
||
# 创建邮件
|
||
msg = MIMEMultipart()
|
||
msg['From'] = SMTP_EMAIL
|
||
msg['To'] = email
|
||
|
||
# 根据用途设置不同的邮件主题和正文
|
||
if purpose == "reset_password":
|
||
msg['Subject'] = "Taiji AI-PAD 密码重置验证码"
|
||
body = f"""
|
||
尊敬的用户：
|
||
|
||
您正在进行密码重置操作，验证码是：{code}
|
||
|
||
验证码有效期为10分钟，请勿泄露给他人。
|
||
|
||
如果您没有进行密码重置操作，请忽略此邮件，您的账户仍然安全。
|
||
|
||
此邮件由系统自动发送，请勿回复。
|
||
|
||
---
|
||
Taiji AI-PAD 团队
|
||
"""
|
||
else:
|
||
msg['Subject'] = "Taiji AI-PAD 注册验证码"
|
||
body = f"""
|
||
尊敬的用户：
|
||
|
||
您的注册验证码是：{code}
|
||
|
||
验证码有效期为10分钟，请勿泄露给他人。
|
||
|
||
如果您没有进行注册操作，请忽略此邮件。
|
||
|
||
此邮件由系统自动发送，请勿回复。
|
||
|
||
---
|
||
Taiji AI-PAD 团队
|
||
"""
|
||
msg.attach(MIMEText(body, 'plain', 'utf-8'))
|
||
|
||
# 发送邮件（使用同步方式，因为 smtplib 不支持异步）
|
||
import asyncio
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(None, _send_email_sync, msg)
|
||
|
||
logger.info("验证码邮件发送成功", email=email)
|
||
return True
|
||
except smtplib.SMTPAuthenticationError as e:
|
||
logger.error(
|
||
"SMTP认证失败，请检查SMTP_PASSWORD是否正确",
|
||
email=email,
|
||
error=str(e),
|
||
smtp_server=SMTP_SERVER,
|
||
smtp_email=SMTP_EMAIL,
|
||
hint="Office365可能需要使用应用专用密码(App Password)而不是普通密码"
|
||
)
|
||
return False
|
||
except smtplib.SMTPException as e:
|
||
logger.error("SMTP发送失败", email=email, error=str(e), smtp_server=SMTP_SERVER)
|
||
return False
|
||
except Exception as e:
|
||
logger.error("验证码邮件发送失败", email=email, error=str(e), error_type=type(e).__name__)
|
||
return False
|
||
|
||
|
||
async def check_rate_limit(email: str) -> tuple[bool, int]:
|
||
"""
|
||
检查验证码发送频率限制
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
|
||
Returns:
|
||
(是否可以发送, 剩余等待秒数)
|
||
"""
|
||
try:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("Redis未连接，跳过频率限制检查")
|
||
return True, 0
|
||
|
||
rate_limit_key = f"verification_rate_limit:{email}"
|
||
|
||
# 检查是否存在频率限制
|
||
ttl = await state.redis_client.ttl(rate_limit_key)
|
||
if ttl > 0:
|
||
logger.warning("验证码发送频率限制", email=email, remaining_seconds=ttl)
|
||
return False, ttl
|
||
|
||
return True, 0
|
||
except Exception as e:
|
||
logger.error("检查频率限制失败", email=email, error=str(e))
|
||
# 出错时允许发送，避免阻塞用户
|
||
return True, 0
|
||
|
||
|
||
async def set_rate_limit(email: str) -> bool:
|
||
"""
|
||
设置验证码发送频率限制
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
|
||
Returns:
|
||
是否设置成功
|
||
"""
|
||
try:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
return False
|
||
|
||
rate_limit_key = f"verification_rate_limit:{email}"
|
||
await state.redis_client.setex(
|
||
rate_limit_key,
|
||
VERIFICATION_CODE_RATE_LIMIT_SECONDS,
|
||
"1"
|
||
)
|
||
return True
|
||
except Exception as e:
|
||
logger.error("设置频率限制失败", email=email, error=str(e))
|
||
return False
|
||
|
||
|
||
async def store_verification_code(email: str, code: str) -> bool:
|
||
"""
|
||
存储验证码到Redis
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
code: 验证码
|
||
|
||
Returns:
|
||
是否存储成功
|
||
"""
|
||
try:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("Redis未连接，无法存储验证码")
|
||
return False
|
||
|
||
key = f"verification_code:{email}"
|
||
await state.redis_client.setex(
|
||
key,
|
||
VERIFICATION_CODE_EXPIRE_SECONDS,
|
||
code
|
||
)
|
||
logger.info("验证码已存储", email=email)
|
||
return True
|
||
except Exception as e:
|
||
logger.error("验证码存储失败", email=email, error=str(e))
|
||
return False
|
||
|
||
|
||
async def verify_code(email: str, code: str) -> bool:
|
||
"""
|
||
验证验证码
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
code: 验证码
|
||
|
||
Returns:
|
||
是否验证成功
|
||
"""
|
||
try:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("Redis未连接，无法验证验证码")
|
||
return False
|
||
|
||
key = f"verification_code:{email}"
|
||
|
||
# 处理Redis集群的MOVED重定向
|
||
max_retries = 3
|
||
for attempt in range(max_retries):
|
||
try:
|
||
stored_code = await state.redis_client.get(key)
|
||
break
|
||
except Exception as e:
|
||
error_str = str(e)
|
||
if "MOVED" in error_str and attempt < max_retries - 1:
|
||
# Redis集群重定向，等待后重试
|
||
import asyncio
|
||
await asyncio.sleep(0.1)
|
||
logger.debug("Redis集群重定向，重试中", attempt=attempt+1, error=error_str)
|
||
continue
|
||
else:
|
||
raise
|
||
|
||
if not stored_code:
|
||
logger.warning("验证码不存在或已过期", email=email, provided_code=code)
|
||
return False
|
||
|
||
# 确保都是字符串类型进行比较
|
||
stored_code = str(stored_code).strip()
|
||
code = str(code).strip()
|
||
|
||
if stored_code != code:
|
||
logger.warning(
|
||
"验证码错误",
|
||
email=email,
|
||
provided_code=code,
|
||
stored_code=stored_code,
|
||
provided_type=type(code).__name__,
|
||
stored_type=type(stored_code).__name__
|
||
)
|
||
return False
|
||
|
||
# 验证成功后删除验证码（同样处理集群重定向）
|
||
for attempt in range(max_retries):
|
||
try:
|
||
await state.redis_client.delete(key)
|
||
break
|
||
except Exception as e:
|
||
error_str = str(e)
|
||
if "MOVED" in error_str and attempt < max_retries - 1:
|
||
import asyncio
|
||
await asyncio.sleep(0.1)
|
||
continue
|
||
else:
|
||
logger.warning("删除验证码失败，但验证已成功", email=email, error=error_str)
|
||
break
|
||
|
||
logger.info("验证码验证成功", email=email)
|
||
return True
|
||
except Exception as e:
|
||
logger.error("验证码验证失败", email=email, code=code, error=str(e), error_type=type(e).__name__)
|
||
return False
|
||
|
||
|
||
async def peek_verification_code(email: str, code: str) -> bool:
|
||
"""检查验证码是否有效，**不删除**（与 verify_code 唯一区别）。
|
||
|
||
用途：注册等多步事务里，先 peek 验证码是否对，等所有 DB / 外部副作用都
|
||
成功后再调一次 verify_code 真消费。这样如果中途失败，用户的验证码不会被
|
||
白白烧掉。
|
||
|
||
Args:
|
||
email: 邮箱
|
||
code: 验证码
|
||
|
||
Returns:
|
||
True 表示验证码存在且匹配；False 表示不存在 / 过期 / 错码 / Redis 不可用
|
||
"""
|
||
try:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("Redis未连接，无法验证验证码")
|
||
return False
|
||
|
||
key = f"verification_code:{email}"
|
||
max_retries = 3
|
||
stored_code = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
stored_code = await state.redis_client.get(key)
|
||
break
|
||
except Exception as e:
|
||
error_str = str(e)
|
||
if "MOVED" in error_str and attempt < max_retries - 1:
|
||
import asyncio
|
||
await asyncio.sleep(0.1)
|
||
continue
|
||
else:
|
||
raise
|
||
|
||
if not stored_code:
|
||
logger.warning("验证码不存在或已过期（peek）", email=email)
|
||
return False
|
||
|
||
stored_code = str(stored_code).strip()
|
||
code = str(code).strip()
|
||
if stored_code != code:
|
||
logger.warning("验证码错误（peek）", email=email)
|
||
return False
|
||
|
||
return True
|
||
except Exception as e:
|
||
logger.error("验证码 peek 失败", email=email, error=str(e),
|
||
error_type=type(e).__name__)
|
||
return False
|
||
|
||
|
||
async def send_and_store_verification_code(email: str, purpose: str = "register") -> Optional[str]:
|
||
"""
|
||
生成、发送并存储验证码
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
purpose: 用途，"register" 表示注册，"reset_password" 表示重置密码
|
||
|
||
Returns:
|
||
验证码（如果成功），None（如果失败）
|
||
"""
|
||
import os
|
||
code = generate_verification_code()
|
||
|
||
# 先存储验证码（即使邮件发送失败，验证码也已存储，可以手动查看Redis）
|
||
store_success = await store_verification_code(email, code)
|
||
if not store_success:
|
||
logger.error("验证码存储失败，无法继续", email=email, purpose=purpose)
|
||
return None
|
||
|
||
# 发送邮件（传递 purpose 参数以使用不同的邮件模板）
|
||
send_success = await send_verification_code(email, code, purpose=purpose)
|
||
|
||
# 测试模式：即使邮件发送失败也返回验证码（仅用于开发/测试环境）
|
||
test_mode = os.getenv("ENABLE_TEST_MODE", "false").lower() == "true" or os.getenv("DEBUG", "false").lower() == "true"
|
||
|
||
if not send_success:
|
||
if test_mode:
|
||
# 测试模式：记录验证码到日志（仅测试环境）
|
||
logger.warning(
|
||
"测试模式：邮件发送失败，但验证码已存储到Redis",
|
||
email=email,
|
||
verification_code=code,
|
||
purpose=purpose,
|
||
hint="验证码已存储到Redis，可通过Redis获取或查看日志（仅测试环境）"
|
||
)
|
||
# 设置发送频率限制
|
||
await set_rate_limit(email)
|
||
return code
|
||
else:
|
||
# 生产模式：邮件发送失败则不返回验证码
|
||
logger.error("邮件发送失败，验证码已存储但未发送", email=email, purpose=purpose)
|
||
return None
|
||
|
||
# 发送成功后设置频率限制
|
||
await set_rate_limit(email)
|
||
|
||
logger.info("验证码已发送并存储", email=email, purpose=purpose)
|
||
return code
|
||
|
||
|
||
async def send_password_reset_code(email: str) -> Optional[str]:
|
||
"""
|
||
发送密码重置验证码（便捷函数）
|
||
|
||
Args:
|
||
email: 邮箱地址
|
||
|
||
Returns:
|
||
验证码（如果成功），None（如果失败）
|
||
"""
|
||
return await send_and_store_verification_code(email, purpose="reset_password")
|
||
|