NeteaseUrl/auth_system/mail.py

# -*- coding: utf-8 -*-
"""SMTP 配置、发信与邮箱验证码。

（自 auth_system.py 拆出，行为与原实现一致；DB 访问与 SMTP 发信均已迁移为
原生异步（aiosqlite / aiosmtplib），不再有阻塞调用藏在 to_thread 里。）
"""

import base64
import contextlib
import re
import secrets
import sqlite3
import ssl
import time
from email.header import Header
from email.mime.text import MIMEText
from email.utils import formataddr
from typing import Any

import aiosmtplib

from .db import (
    _q,
    _x,
    get_secret_key,
    write_transaction,
)
from .security import (
    _normalize_email,
)

# ============ 邮箱验证码 / SMTP ============

# 验证码猜错计数（内存态，重启清零可接受）——同邮箱累计错 5 次作废验证码
# 键='purpose:任意邮箱' 无上限可被慢速灌大——超限清最旧
_CODE_FAILS: dict = {}

_CODE_FAILS_MAX = 5000


def _code_fails_trim() -> None:
    """超上限时按写入序丢弃最旧的一半（仅防内存无界，不承担安全语义）"""
    if len(_CODE_FAILS) > _CODE_FAILS_MAX:
        for k in list(_CODE_FAILS.keys())[:len(_CODE_FAILS) // 2]:
            _CODE_FAILS.pop(k, None)


_SMTP_ENC_PREFIX = 'enc1:'   # 信封加密标记前缀（存量明文读取时透明升级）


async def _smtp_password_encrypt(plaintext: str) -> str:
    """用 data/secret.key 派生 Fernet 密钥加密 SMTP 授权码。

    原实现明文落库、并随数据库备份一起导出，等于扩大泄露半径。"""
    import hashlib as _hl

    from cryptography.fernet import Fernet
    key = _hl.sha256((await get_secret_key()).encode('utf-8')).digest()
    f = Fernet(base64.urlsafe_b64encode(key))
    return _SMTP_ENC_PREFIX + f.encrypt(plaintext.encode('utf-8')).decode('utf-8')


async def _smtp_password_decrypt(ciphertext: str) -> str:
    if not ciphertext.startswith(_SMTP_ENC_PREFIX):
        return ciphertext   # 存量明文：透明兼容，下次保存时升级为密文
    import hashlib as _hl

    from cryptography.fernet import Fernet
    key = _hl.sha256((await get_secret_key()).encode('utf-8')).digest()
    f = Fernet(base64.urlsafe_b64encode(key))
    return f.decrypt(ciphertext[len(_SMTP_ENC_PREFIX):].encode('utf-8')).decode('utf-8')


async def get_smtp_config() -> dict[str, Any]:
    rows = await _q('SELECT * FROM smtp_config WHERE id = 1')
    cfg = dict(rows[0]) if rows else {'host': '', 'port': 587, 'sender': '', 'password': ''}
    # 读取时解密密码——调用方（SMTP 连接/后台展示）拿到的始终是明文
    if cfg.get('password'):
        try:
            cfg['password'] = await _smtp_password_decrypt(cfg['password'])
        except Exception:
            cfg['password'] = ''
    return cfg


async def save_smtp_config(host: str, port: int, sender: str, password: str) -> None:
    enc_pwd = await _smtp_password_encrypt(password) if password else ''
    await _x('UPDATE smtp_config SET host = ?, port = ?, sender = ?, password = ?, updated_at = ? WHERE id = 1',
             (host, int(port or 587), sender, enc_pwd, int(time.time())))


async def smtp_configured() -> bool:
    cfg = await get_smtp_config()
    return bool(cfg['host'] and cfg['sender'] and cfg['password'])


async def _smtp_send(cfg: dict[str, Any], sender: str, to: str, raw: str) -> None:
    """原生异步发信（aiosmtplib）。

    465 走隐式 TLS（use_tls），其余端口走显式 STARTTLS（start_tls）——与原先
    smtplib 的 SMTP_SSL / starttls 分支一一对应；两条路都用显式构造的 TLS 上下文
    校验证书与主机名，不依赖各运行时默认值。原实现是 smtplib 同步栈再经
    asyncio.to_thread 卸载，换成 aiosmtplib 后不再占用线程池。
    """
    port = int(cfg['port'] or 587)
    ctx = ssl.create_default_context()
    implicit_tls = port == 465
    smtp = aiosmtplib.SMTP(
        hostname=cfg['host'],
        port=port,
        timeout=15,
        use_tls=implicit_tls,
        start_tls=not implicit_tls,
        tls_context=ctx,
    )
    await smtp.connect()
    try:
        await smtp.login(cfg['sender'], cfg['password'])
        await smtp.sendmail(sender, [to], raw)
    finally:
        # 关连接失败不影响「信已发出」的事实，也绝不能盖掉上面的真实异常
        with contextlib.suppress(Exception):
            await smtp.quit()


async def send_email(to: str, subject: str, body: str) -> None:
    """收件人格式校验下沉到本入口——未来任何新调用方漏校验也不再成为
    CRLF/信封注入点（现有调用点的校验仅作 UX 提示）。"""
    to = _normalize_email(to)
    if not re.match(r'^[\w.%+-]+@[\w.-]+\.[A-Za-z]{2,}$', to, re.ASCII):
        raise ValueError('收件邮箱格式不正确')
    cfg = await get_smtp_config()
    msg = MIMEText(body, 'plain', 'utf-8')
    # email 库推荐写法：Header 对象由 Generator._fold 的 Header 分支做 RFC 2047 编码
    # （typeshed 把 Message.__setitem__ 的值参数收窄成 str，属桩过窄；运行时已实测正确）
    msg['Subject'] = Header(subject, 'utf-8')  # type: ignore[assignment]
    msg['From'] = formataddr(('云村小助手', cfg['sender']))
    msg['To'] = to
    # SMTP 已是原生异步（aiosmtplib），直接 await——不再有同步网络 IO 需要卸载
    await _smtp_send(cfg, cfg['sender'], to, msg.as_string())


_EMAIL_SENT_24H: dict = {}   # 同邮箱 24h 发送计数（内存态，重启清零可接受）


# 验证码用途隔离——注册码不能用于找回密码，反之亦然（库行带 purpose，
# 消费/作废/还原均按用途匹配）。邮件主题与正文按用途区分，用户一眼可辨。
# 新增 admin2fa（管理员两步验证）——发码前置条件是密码已验证正确，
# 故豁免 24h 5 封与 60s 冷却（否则管理员每天最多登录 5 次会被自己锁死）；
# 滥用面由"票据仅密码正确时签发 + 票据尝试上限 + 登录限流"约束
_CODE_PURPOSES = ('register', 'reset', 'admin2fa')

_CODE_EMAIL_TEMPLATES = {
    'register': ('云村小助手 注册验证码',
                 '您正在注册云村小助手账号。\n\n'
                 '验证码：{code}\n\n'
                 '验证码 10 分钟内有效，请勿泄露给他人。\n'
                 '如非本人操作，请忽略本邮件。'),
    'reset': ('云村小助手 找回密码验证码',
              '您正在找回云村小助手账号的登录密码。\n\n'
              '验证码：{code}\n\n'
              '验证码 10 分钟内有效，仅用于重置密码，请勿泄露给他人。\n'
              '如非本人操作，请忽略本邮件（您的密码不会被修改）。'),
    'admin2fa': ('云村小助手 管理员登录验证码',
                 '您正在登录云村小助手管理员账号（两步验证第二步）。\n\n'
                 '验证码：{code}\n\n'
                 '验证码 10 分钟内有效，请勿泄露给任何人。\n'
                 '如非本人操作，请立即修改管理员密码并检查邮箱安全。'),
}


async def send_verify_code(email: str, purpose: str = 'register') -> None:
    """生成并发送邮箱验证码（10 分钟有效）（先发信成功再落库；同用途旧码即时作废）。
① 同邮箱 24h 硬上限 5 封（防邮件轰炸代理/发信配额消耗）；
    ② email 先归一化（Gmail +tag 别名剥离）——别名刷号被邮箱查重堵死。
    purpose 区分注册/找回（文案不同 + 库行用途隔离）；24h 计数移到
    发信成功之后（原实现 SMTP 失败也烧掉额度）。"""
    if purpose not in _CODE_PURPOSES:
        raise ValueError('验证码用途不正确')
    email = _normalize_email(email or '')
    # re.ASCII——\w 的 Unicode 语义会放行"中文邮箱@例子.com"这类
    # 非 ASCII 地址（SMTP 阶段才失败，路由兜底成 500）；加 254 长度上限（RFC 口径，
    # 原实现 5000 字符邮箱可入库）
    if len(email) > 254 or not re.match(r'^[\w.%+-]+@[\w.-]+\.[A-Za-z]{2,}$', email, re.ASCII):
        raise ValueError('邮箱格式不正确')
    if not await smtp_configured():
        raise RuntimeError('邮箱服务未配置，请联系管理员')
    # 同邮箱 24h ≤5 封（内存计数，不再依赖会被消费 DELETE 清零的 DB 行数）；
    # admin2fa 豁免 24h 上限与 60s 冷却（见 _CODE_PURPOSES 处注释）
    now_ts = time.time()
    sent = [t for t in _EMAIL_SENT_24H.get(email, []) if now_ts - t < 86400]
    if purpose != 'admin2fa' and len(sent) >= 5:
        raise RuntimeError('该邮箱验证码发送已达上限（24 小时 5 封），请明日再试')
    # 60 秒内限制重发（按同邮箱同用途的最近一封计）
    if purpose != 'admin2fa':
        rows = await _q('SELECT created_at FROM email_codes WHERE email = ? AND purpose = ? ORDER BY id DESC LIMIT 1',
                        (email, purpose))
        if rows and time.time() - rows[0]['created_at'] < 60:
            raise RuntimeError('发送过于频繁，请 1 分钟后再试')
    code = f'{secrets.randbelow(1000000):06d}'
    subject, body_tpl = _CODE_EMAIL_TEMPLATES[purpose]
    # 先发送：SMTP 失败不留任何有效验证码与冷却记录；成功后才计入 24h 额度（admin2fa 不计数）
    await send_email(email, subject, body_tpl.format(code=code))
    if purpose != 'admin2fa':
        sent.append(now_ts)
        _EMAIL_SENT_24H[email] = sent
        # 键只增不减（旧时间戳被过滤但空键留存）——顺手剔除过期
        # 空键并做容量上限，防慢速灌任意格式邮箱造成无界增长
        if len(_EMAIL_SENT_24H) > 5000:
            for k in [k for k, v in _EMAIL_SENT_24H.items() if not v][:2500]:
                _EMAIL_SENT_24H.pop(k, None)
            if len(_EMAIL_SENT_24H) > 5000:
                for k in list(_EMAIL_SENT_24H.keys())[:len(_EMAIL_SENT_24H) // 2]:
                    _EMAIL_SENT_24H.pop(k, None)
    async with write_transaction() as conn:
        # 新码生效同时作废同邮箱同用途旧码；顺手清理全表过期码
        await conn.execute('DELETE FROM email_codes WHERE (email = ? AND purpose = ?) OR expires_at <= ?',
                           (email, purpose, int(time.time())))
        await conn.execute('INSERT INTO email_codes (email, code, expires_at, created_at, purpose) VALUES (?, ?, ?, ?, ?)',
                           (email, code, int(time.time() + 600), int(time.time()), purpose))


async def check_verify_code(email: str, code: str, purpose: str = 'register') -> int:
    """校验验证码（单条 DELETE 原子消费，rowcount=1 才算通过，杜绝并发双读）。
    成功返回剩余有效秒数（真值），失败返回 0。
    同一邮箱累计错 5 次→该用途验证码作废（防 10 分钟窗口内挂机猜码）。
    WHERE 带 purpose——注册码与找回码互不可用。"""
    if purpose not in _CODE_PURPOSES:
        return 0
    email = _normalize_email(email or '')
    fail_key = purpose + ':' + email   # 失败计数按"用途:邮箱"隔离，互不连带
    row = None
    try:
        async with write_transaction() as conn:
            cur = await conn.execute('DELETE FROM email_codes WHERE email = ? AND code = ? AND purpose = ? AND expires_at > ? '
                                     'RETURNING expires_at',
                                     (email, (code or '').strip(), purpose, int(time.time())))
            row = await cur.fetchone()
    except sqlite3.OperationalError as _oe:
        # 老版本 SQLite 无 RETURNING：回退为删除后判 rowcount。
        # 仅"不支持 RETURNING"走回退——库忙（database is locked）原样
        # 上抛，否则第一次等满 busy 15s、回退再等 15s，全程持锁冻结全站写路径
        if 'RETURNING' not in str(_oe) and 'syntax' not in str(_oe).lower():
            raise
        async with write_transaction() as conn:
            cur = await conn.execute('DELETE FROM email_codes WHERE email = ? AND code = ? AND purpose = ? AND expires_at > ?',
                                     (email, (code or '').strip(), purpose, int(time.time())))
            if cur.rowcount == 1:
                _CODE_FAILS.pop(fail_key, None)
                return 60
        row = None
    if row:
        _CODE_FAILS.pop(fail_key, None)
        return max(0, int(row[0]) - int(time.time()))
    # 猜错：失败计数 +1，累计 5 次直接作废该用途的验证码
    n = _CODE_FAILS.get(fail_key, 0) + 1
    _CODE_FAILS[fail_key] = n
    _code_fails_trim()
    if n >= 5:
        _CODE_FAILS.pop(fail_key, None)
        try:
            await _x('DELETE FROM email_codes WHERE email = ? AND purpose = ?', (email, purpose))
        except Exception:
            pass
    return 0


async def restore_verify_code(email: str, code: str, ttl_seconds: int, purpose: str = 'register') -> None:
    """注册流程后续步骤失败时还原已消费的验证码（第二轮：建号失败不应吞掉验证码）"""
    if ttl_seconds <= 0:
        return
    if purpose not in _CODE_PURPOSES:
        return
    email = _normalize_email(email or '')
    code = (code or '').strip()
    if not email or not code:
        return
    await _x('INSERT INTO email_codes (email, code, expires_at, created_at, purpose) VALUES (?, ?, ?, ?, ?)',
             (email, code, int(time.time() + ttl_seconds), int(time.time()), purpose))