document-management/backend/dms/services/captcha_service.py

"""登录图片验证码：内存签发/校验 + PyMuPDF 渲染 PNG（安全需求4）。

- 单实例部署（launcher 单进程）下内存存储足够；
- 4 位字符（剔除易混淆 0O1I），120 秒过期；
- 校验一次性消费（无论成败），防止同一验证码反复试密码；
- 校验忽略大小写与首尾空白。
"""

from __future__ import annotations

import base64
import secrets
import threading
import time
import uuid
import random

from dms.common.errors import InvalidCaptchaError

_CAPTCHA_CHARS = "23456789ABCDEFGHJKLMNPQRSTUVWXYZ"
_CODE_LENGTH = 4
_EXPIRES_SECONDS = 120
_MAX_ENTRIES = 1000

_lock = threading.Lock()
# captcha_id -> (code 大写归一, 过期时刻 monotonic)
_entries: dict[str, tuple[str, float]] = {}


def issue() -> dict[str, str]:
    """签发新验证码，返回 ``{"captchaId": ..., "image": "data:image/png;base64,..."}``。"""
    code = "".join(secrets.choice(_CAPTCHA_CHARS) for _ in range(_CODE_LENGTH))
    captcha_id = uuid.uuid4().hex
    now = time.monotonic()
    with _lock:
        _purge_expired(now)
        if len(_entries) >= _MAX_ENTRIES:
            # 上限保护：极端刷量时整体清空（存量本就即将过期，不影响正常用户）
            _entries.clear()
        _entries[captcha_id] = (code.upper(), now + _EXPIRES_SECONDS)
    return {
        "captchaId": captcha_id,
        "image": "data:image/png;base64,"
        + base64.b64encode(_render_png(code)).decode("ascii"),
    }


def verify(captcha_id: object, code: object) -> None:
    """校验并消费验证码；缺失/错误/过期均抛 :class:`InvalidCaptchaError`。"""
    if not isinstance(captcha_id, str) or not captcha_id:
        raise InvalidCaptchaError()
    if not isinstance(code, str) or not code.strip():
        raise InvalidCaptchaError()
    with _lock:
        entry = _entries.pop(captcha_id, None)
    if entry is None or entry[1] <= time.monotonic():
        raise InvalidCaptchaError()
    if entry[0] != code.strip().upper():
        raise InvalidCaptchaError()


def _purge_expired(now: float) -> None:
    for key in [k for k, (_, expires) in _entries.items() if expires <= now]:
        _entries.pop(key, None)


def _render_png(code: str) -> bytes:
    """用 PyMuPDF（项目既有依赖）绘制验证码位图：随机颜色/基线/字号 + 干扰线。"""
    import fitz

    width, height = 132, 44
    rng = random.Random()
    doc = fitz.open()
    try:
        page = doc.new_page(width=width, height=height)
        page.draw_rect(page.rect, color=None, fill=(0.96, 0.98, 1.0))
        for _ in range(4):
            page.draw_line(
                fitz.Point(rng.uniform(0, width), rng.uniform(0, height)),
                fitz.Point(rng.uniform(0, width), rng.uniform(0, height)),
                color=(rng.random() * 0.6, rng.random() * 0.6, rng.random() * 0.6),
                width=1,
            )
        x = 14.0
        for ch in code:
            page.insert_text(
                fitz.Point(x, rng.uniform(28, 38)),
                ch,
                fontname="hebo",
                fontsize=rng.uniform(22, 26),
                color=(
                    rng.uniform(0, 0.45),
                    rng.uniform(0, 0.45),
                    rng.uniform(0, 0.45),
                ),
            )
            x += 27
        return page.get_pixmap().tobytes("png")
    finally:
        doc.close()