document-management/backend/tests/dms/test_captcha.py

"""登录图片验证码（安全需求4）：

- 签发结构：captchaId + data URL PNG，明文不外泄；
- 校验：成功、忽略大小写、一次性消费（无论成败）、过期、错误/缺失；
- 接口层：开启时 GET /auth/captcha 返回图片与 enabled=true；
  登录缺验证码字段或校验失败在查库前即被拒（INVALID_CAPTCHA，不计锁定）。

不依赖测试库：全部在 config 覆盖 + 内存 store 上完成。
"""

from __future__ import annotations

import base64
import time

import pytest

from dms.common.errors import InvalidCaptchaError
from dms.services import captcha_service as cs


@pytest.fixture()
def clean_store():
    cs._entries.clear()
    yield
    cs._entries.clear()


def _issued_with_plain_code() -> tuple[str, str]:
    data = cs.issue()
    return data["captchaId"], cs._entries[data["captchaId"]][0]


def test_issue_shape(clean_store):
    """签发返回 captchaId 与 data:image/png;base64 前缀图片，且为有效 PNG 头。"""
    data = cs.issue()
    assert set(data) == {"captchaId", "image"}
    assert data["captchaId"]
    prefix = "data:image/png;base64,"
    assert data["image"].startswith(prefix)
    raw = base64.b64decode(data["image"][len(prefix):])
    assert raw.startswith(b"\x89PNG")


def test_verify_success_and_case_insensitive(clean_store):
    captcha_id, code = _issued_with_plain_code()
    cs.verify(captcha_id, code.lower())


def test_verify_consumes_once(clean_store):
    """一次性消费：第二次校验（即使正确）也拒绝，防重放。"""
    captcha_id, code = _issued_with_plain_code()
    cs.verify(captcha_id, code)
    with pytest.raises(InvalidCaptchaError):
        cs.verify(captcha_id, code)


def test_verify_wrong_code_fails_and_consumes(clean_store):
    """错误尝试即作废，防止同一验证码反复试密码。"""
    captcha_id, code = _issued_with_plain_code()
    with pytest.raises(InvalidCaptchaError):
        cs.verify(captcha_id, "X" * len(code))
    with pytest.raises(InvalidCaptchaError):
        cs.verify(captcha_id, code)


def test_verify_expired(clean_store):
    captcha_id, code = _issued_with_plain_code()
    # 把过期时刻拉回过去（monotonic 时钟，120 秒窗口）
    cs._entries[captcha_id] = (code, time.monotonic() - 1)
    with pytest.raises(InvalidCaptchaError):
        cs.verify(captcha_id, code)


def test_verify_missing_arguments(clean_store):
    with pytest.raises(InvalidCaptchaError):
        cs.verify(None, None)
    with pytest.raises(InvalidCaptchaError):
        cs.verify("whatever", "  ")


def test_captcha_endpoint_and_login_guard(flask_app):
    """/auth/captcha 开关行为 + 开启时登录请求体扩为五键、校验先于查库。"""
    client = flask_app.test_client()

    # 开关关闭（conftest 设定）：仅返回 enabled=false，登录不再被验证码拦截
    response = client.get("/api/v1/auth/captcha")
    assert response.status_code == 200
    assert response.get_json()["data"] == {"enabled": False}
    three_keys = {"username": "u", "password": "p", "keepSignedIn": False}

    flask_app.config["DMS_CAPTCHA_ENABLED"] = True
    cs._entries.clear()
    try:
        response = client.get("/api/v1/auth/captcha")
        assert response.status_code == 200
        payload = response.get_json()["data"]
        assert payload["enabled"] is True
        assert set(payload) == {"enabled", "captchaId", "image"}

        # 缺 captchaId/captchaCode：结构校验 400 INVALID_ARGUMENT
        response = client.post("/api/v1/auth/login", json=three_keys)
        assert response.status_code == 400
        assert response.get_json()["code"] == "INVALID_ARGUMENT"

        # 带未知 captchaId：查库前即 400 INVALID_CAPTCHA（无测试库依赖）
        response = client.post(
            "/api/v1/auth/login",
            json={
                **three_keys,
                "captchaId": "not-exist",
                "captchaCode": "ABCD",
            },
        )
        assert response.status_code == 400
        assert response.get_json()["code"] == "INVALID_CAPTCHA"

        # 真实签发的验证码先被消费（模拟已使用过），重放同样拒绝
        captcha_id = payload["captchaId"]
        code = cs._entries[captcha_id][0]
        cs.verify(captcha_id, code)
        response = client.post(
            "/api/v1/auth/login",
            json={
                **three_keys,
                "captchaId": captcha_id,
                "captchaCode": code,
            },
        )
        assert response.status_code == 400
        assert response.get_json()["code"] == "INVALID_CAPTCHA"
    finally:
        flask_app.config["DMS_CAPTCHA_ENABLED"] = False
        cs._entries.clear()