forked from xiaohei/taiji-AI-PAD
324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""
|
||
Heicode magic-link 邮箱登录 —— Redis 一次性 token/code 存取 + 登录链接邮件发送
|
||
|
||
本模块为 Heicode magic-link 登录新增的独立支撑层，契约见
|
||
`Docs/Heicode-magic-link邮箱登录-给mcp-server的对接需求.md`（§2/§7/§9/§11/§12）。
|
||
|
||
设计要点（与既有逻辑零耦合，绝不改动密码登录/注册验证码链路）：
|
||
- **复用** `email_verification` 的 SMTP 通道（smtp.189.cn / taijiagent@189.cn）发送
|
||
登录链接邮件；
|
||
- **复用** `state.redis_client` 存一次性凭证，沿用现有验证码同款 one-time 模式
|
||
（`setex` 落地 + 命中即 `delete`，含 Redis 集群 MOVED 重定向重试）；
|
||
- token / code 均为短 TTL、一次性：
|
||
- magic-link token：TTL 600s（邮件链接里的凭证）
|
||
- 一次性 code：TTL 120s（landing 校验通过后换取登录态的凭证）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import secrets
|
||
from typing import Optional
|
||
|
||
import structlog
|
||
|
||
from config import settings
|
||
from app.state import get_state
|
||
# 复用现有 SMTP 通道与同步发信实现，不另起炉灶
|
||
from app.email_verification import (
|
||
SMTP_EMAIL,
|
||
SMTP_PASSWORD,
|
||
_send_email_sync,
|
||
)
|
||
from email.mime.text import MIMEText
|
||
from email.mime.multipart import MIMEMultipart
|
||
|
||
logger = structlog.get_logger(__name__)
|
||
|
||
# ===== TTL / Redis key 约定（契约 §2 / §7.1）=====
|
||
MAGIC_LINK_TOKEN_TTL_SECONDS = 600 # 邮件链接 token：10 分钟
|
||
MAGIC_LINK_CODE_TTL_SECONDS = 120 # 一次性 code：≤2 分钟
|
||
|
||
_TOKEN_KEY_PREFIX = "magic_link_token:"
|
||
_CODE_KEY_PREFIX = "magic_link_code:"
|
||
# 邮箱维度限流（§13.2「邮箱限流」）。**独立**于注册/忘记密码共用的
|
||
# verification_rate_limit:{email}，避免 magic-link 与那两条流程互相误伤。
|
||
_EMAIL_RATE_LIMIT_KEY_PREFIX = "magic_link_rate_limit:"
|
||
EMAIL_RATE_LIMIT_SECONDS = 60 # 同一邮箱 60s 内只接受一次申请
|
||
|
||
_MAX_REDIS_RETRIES = 3
|
||
|
||
|
||
def generate_magic_link_token() -> str:
|
||
"""生成 magic-link token（URL-safe，随邮件链接下发）。"""
|
||
return secrets.token_urlsafe(32)
|
||
|
||
|
||
def generate_one_time_code() -> str:
|
||
"""生成一次性 code（URL-safe，landing 302 回跳给客户端）。"""
|
||
return secrets.token_urlsafe(24)
|
||
|
||
|
||
async def _redis_get(key: str) -> Optional[str]:
|
||
"""带 Redis 集群 MOVED 重定向重试的 GET（与 email_verification 同款）。"""
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("magic_link_redis_unavailable", op="get")
|
||
return None
|
||
for attempt in range(_MAX_REDIS_RETRIES):
|
||
try:
|
||
return await state.redis_client.get(key)
|
||
except Exception as e: # noqa: BLE001
|
||
if "MOVED" in str(e) and attempt < _MAX_REDIS_RETRIES - 1:
|
||
import asyncio
|
||
await asyncio.sleep(0.1)
|
||
continue
|
||
raise
|
||
return None
|
||
|
||
|
||
async def _redis_delete(key: str) -> None:
|
||
"""带 MOVED 重试的 DELETE；删除失败不致命（一次性消费已凭 get 判定）。"""
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
return
|
||
for attempt in range(_MAX_REDIS_RETRIES):
|
||
try:
|
||
await state.redis_client.delete(key)
|
||
return
|
||
except Exception as e: # noqa: BLE001
|
||
if "MOVED" in str(e) and attempt < _MAX_REDIS_RETRIES - 1:
|
||
import asyncio
|
||
await asyncio.sleep(0.1)
|
||
continue
|
||
logger.warning("magic_link_redis_delete_failed", key_prefix=key[:24], error=str(e))
|
||
return
|
||
|
||
|
||
async def _setex(key: str, ttl: int, value: str) -> bool:
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
logger.warning("magic_link_redis_unavailable", op="setex")
|
||
return False
|
||
try:
|
||
await state.redis_client.setex(key, ttl, value)
|
||
return True
|
||
except Exception as e: # noqa: BLE001
|
||
logger.error("magic_link_redis_setex_failed", key_prefix=key[:24], error=str(e))
|
||
return False
|
||
|
||
|
||
# ===== 邮箱维度限流（§13.2）=====
|
||
|
||
async def check_email_rate_limit(email: str) -> tuple[bool, int]:
|
||
"""检查同一邮箱的申请频率限制。
|
||
|
||
Returns: (是否可以申请, 剩余等待秒数)。Redis 不可用时放行（不阻塞用户），
|
||
与 email_verification.check_rate_limit 的容错口径一致。
|
||
"""
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
return True, 0
|
||
try:
|
||
ttl = await state.redis_client.ttl(_EMAIL_RATE_LIMIT_KEY_PREFIX + email)
|
||
if ttl and ttl > 0:
|
||
return False, ttl
|
||
return True, 0
|
||
except Exception as e: # noqa: BLE001
|
||
logger.error("magic_link_email_rate_limit_check_failed", email=email, error=str(e))
|
||
return True, 0
|
||
|
||
|
||
async def set_email_rate_limit(email: str) -> None:
|
||
"""对该邮箱设置 60s 申请冷却。**对存在/不存在的邮箱一律设置**——否则
|
||
「存在→后续 429 / 不存在→后续 200」会泄漏邮箱存在性，破坏防枚举（D-2）。"""
|
||
state = get_state()
|
||
if not state.redis_client:
|
||
return
|
||
try:
|
||
await state.redis_client.setex(
|
||
_EMAIL_RATE_LIMIT_KEY_PREFIX + email, EMAIL_RATE_LIMIT_SECONDS, "1"
|
||
)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.error("magic_link_email_rate_limit_set_failed", email=email, error=str(e))
|
||
|
||
|
||
# ===== token：email + state 绑定 =====
|
||
|
||
async def store_magic_link_token(
|
||
token: str,
|
||
email: str,
|
||
state: str,
|
||
web_callback: Optional[str] = None,
|
||
accept_url: Optional[str] = None,
|
||
ttl: int = MAGIC_LINK_TOKEN_TTL_SECONDS,
|
||
) -> bool:
|
||
"""存 magic-link token → {email, state[, web_callback, accept_url]}，TTL 默认 600s。
|
||
|
||
- 桌面登录：只传 email/state（landing 302 到 heicode://）。
|
||
- 企业邀请 web 模式（Q2-B）：额外传 web_callback + accept_url（landing 302 到
|
||
web_callback?code=&state=&redirect=accept_url），并可用更长 ttl（邀请信可能隔时点）。
|
||
"""
|
||
data = {"email": email, "state": state}
|
||
if web_callback:
|
||
data["web_callback"] = web_callback
|
||
if accept_url:
|
||
data["accept_url"] = accept_url
|
||
return await _setex(_TOKEN_KEY_PREFIX + token, ttl, json.dumps(data))
|
||
|
||
|
||
async def consume_magic_link_token(token: str) -> Optional[dict]:
|
||
"""一次性消费 token：命中则返回 {email, state} 并删除；否则 None。"""
|
||
if not token:
|
||
return None
|
||
key = _TOKEN_KEY_PREFIX + token
|
||
raw = await _redis_get(key)
|
||
if not raw:
|
||
return None
|
||
await _redis_delete(key)
|
||
try:
|
||
return json.loads(raw)
|
||
except (ValueError, TypeError):
|
||
logger.warning("magic_link_token_payload_corrupt")
|
||
return None
|
||
|
||
|
||
# ===== code：user_id + email + state 绑定 =====
|
||
|
||
async def store_one_time_code(code: str, user_id: str, email: str, state: str) -> bool:
|
||
"""存一次性 code → {user_id, email, state}，TTL ≤120s。"""
|
||
payload = json.dumps({"user_id": user_id, "email": email, "state": state})
|
||
return await _setex(_CODE_KEY_PREFIX + code, MAGIC_LINK_CODE_TTL_SECONDS, payload)
|
||
|
||
|
||
async def consume_one_time_code(code: str) -> Optional[dict]:
|
||
"""一次性消费 code：命中则返回 {user_id, email, state} 并删除；否则 None。"""
|
||
if not code:
|
||
return None
|
||
key = _CODE_KEY_PREFIX + code
|
||
raw = await _redis_get(key)
|
||
if not raw:
|
||
return None
|
||
await _redis_delete(key)
|
||
try:
|
||
return json.loads(raw)
|
||
except (ValueError, TypeError):
|
||
logger.warning("magic_link_code_payload_corrupt")
|
||
return None
|
||
|
||
|
||
# ===== 登录链接邮件（复用现有 SMTP 通道）=====
|
||
|
||
async def send_magic_link_email(email: str, link_url: str) -> bool:
|
||
"""发送 magic-link 登录链接邮件。
|
||
|
||
复用 email_verification 的 SMTP 通道。两道闸：
|
||
1. `MAGIC_LINK_EMAIL_ENABLED`（默认 False）—— 契约 §13.2「D-5 签字前默认不外发」，
|
||
未置 true 一律 mock（不外发），即便 SMTP 可用也不发，防止端点被滥用。
|
||
2. `SMTP_PASSWORD` 未注入 —— 通道不可用。
|
||
任一不满足都返回 False、不抛异常，让上层 `request` 端点照常返回 200（防枚举）。
|
||
敏感链接（含 token）仅在 DEBUG 下打印，避免普通日志留存凭证（§1.4）。
|
||
"""
|
||
import os
|
||
debug = os.getenv("DEBUG", "false").lower() == "true" or \
|
||
os.getenv("ENABLE_TEST_MODE", "false").lower() == "true"
|
||
|
||
# 闸 1：未签字 → mock 不外发
|
||
if not settings.magic_link_email_enabled:
|
||
if debug:
|
||
logger.warning("magic_link_email_mock", email=email, link_url=link_url,
|
||
hint="MAGIC_LINK_EMAIL_ENABLED=false → mock，不外发（§13.2）")
|
||
else:
|
||
logger.info("magic_link_email_mock", email=email,
|
||
hint="MAGIC_LINK_EMAIL_ENABLED=false → mock，不外发（链接仅 DEBUG 打印）")
|
||
return False
|
||
|
||
# 闸 2：通道不可用
|
||
if not SMTP_PASSWORD:
|
||
logger.error("magic_link_email_smtp_unconfigured", email=email)
|
||
return False
|
||
|
||
try:
|
||
msg = MIMEMultipart()
|
||
msg["From"] = SMTP_EMAIL
|
||
msg["To"] = email
|
||
msg["Subject"] = "Taiji AI-PAD 登录链接"
|
||
body = f"""
|
||
尊敬的用户：
|
||
|
||
您正在登录 HeiCode 客户端。请在 10 分钟内，**在已安装 HeiCode 的同一台设备上**
|
||
点击下面的链接完成登录：
|
||
|
||
{link_url}
|
||
|
||
提示：此链接仅用于本次登录，点击后会自动回跳到本机 HeiCode 客户端。
|
||
请务必在安装了 HeiCode 的同一台设备上打开此链接，否则客户端无法收到回跳。
|
||
|
||
如果您没有发起登录，请忽略此邮件，您的账户仍然安全。
|
||
|
||
此邮件由系统自动发送，请勿回复。
|
||
|
||
---
|
||
Taiji AI-PAD 团队
|
||
"""
|
||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||
|
||
import asyncio
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(None, _send_email_sync, msg)
|
||
logger.info("magic_link_email_sent", email=email)
|
||
return True
|
||
except Exception as e: # noqa: BLE001
|
||
logger.error("magic_link_email_send_failed", email=email, error=str(e),
|
||
error_type=type(e).__name__)
|
||
return False
|
||
|
||
|
||
async def send_org_invite_email(email: str, accept_link: str, org_name: Optional[str] = None) -> bool:
|
||
"""发送企业邀请信（Q2-B）。链接指向 mcp landing（web 模式），点击后经 HM 回调建会话并入组。
|
||
|
||
与 magic-link 同一发信总闸（`MAGIC_LINK_EMAIL_ENABLED`）+ SMTP 通道。任一不满足返回 False、不抛异常。
|
||
"""
|
||
if not settings.magic_link_email_enabled:
|
||
logger.info("org_invite_email_mock", email=email,
|
||
hint="MAGIC_LINK_EMAIL_ENABLED=false → mock，不外发")
|
||
return False
|
||
if not SMTP_PASSWORD:
|
||
logger.error("org_invite_email_smtp_unconfigured", email=email)
|
||
return False
|
||
|
||
org = (org_name or "").strip()
|
||
org_line = f"您被邀请加入组织 **{org}**。" if org else "您收到一封企业邀请。"
|
||
try:
|
||
msg = MIMEMultipart()
|
||
msg["From"] = SMTP_EMAIL
|
||
msg["To"] = email
|
||
msg["Subject"] = (f"Heicode 企业邀请 - {org}" if org else "Heicode 企业邀请")
|
||
body = f"""
|
||
您好：
|
||
|
||
{org_line}
|
||
请点击下面的链接完成登录并接受邀请：
|
||
|
||
{accept_link}
|
||
|
||
点击后会自动完成身份验证并跳转到接受邀请页面。链接有时效，请尽快点击。
|
||
|
||
如果您并未预期收到此邀请，请忽略此邮件。
|
||
|
||
此邮件由系统自动发送，请勿回复。
|
||
|
||
---
|
||
Heicode 团队
|
||
"""
|
||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||
import asyncio
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(None, _send_email_sync, msg)
|
||
logger.info("org_invite_email_sent", email=email, org=org or None)
|
||
return True
|
||
except Exception as e: # noqa: BLE001
|
||
logger.error("org_invite_email_send_failed", email=email, error=str(e),
|
||
error_type=type(e).__name__)
|
||
return False
|