llm-in-text/backend/captcha_api.py

"""
验证码和 Cookie 策略管理模块

提供功能:
1. 图形验证码生成与验证 API
2. 现代 Cookie 策略管理 (HttpOnly, Secure, SameSite)
3. 验证码结果持久化到 Cookie
"""

import random
import string
import json
from typing import Optional
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi.responses import JSONResponse

router = APIRouter(prefix="/captcha", tags=["验证码"])


# ==================== 数据模型 ====================

class CaptchaConfig:
    """验证码配置"""
    LENGTH = 6  # 验证码长度
    CHARSET = string.ascii_letters + string.digits  # 字符集: 大小写字母+数字
    EXPIRE_SECONDS = 3600  # 过期时间: 1小时


class CaptchaResult:
    """验证码结果"""
    def __init__(self, text: str):
        self.text = text
        self.created_at = int(__import__('time').time())
    
    @property
    def is_expired(self) -> bool:
        now = int(__import__('time').time())
        return (now - self.created_at) > CaptchaConfig.EXPIRE_SECONDS


# ==================== 全局状态 ====================

# 内存中的验证码存储 (生产环境建议用 Redis)
_active_captchas: dict[str, CaptchaResult] = {}


# ==================== 验证码 API ====================

@router.get("/generate", summary="生成新验证码")
async def generate_captcha(
    response: Response,
    use_cookie: bool = False,  # 是否通过 Cookie 传递验证码文本
    length: int = CaptchaConfig.LENGTH,
):
    """
    生成新的验证码
    
    - **use_cookie**: 是否同时设置 Cookie (方便前端读取)
    - **length**: 验证码长度 (4-10)
    
    返回:
    - **request_id**: 验证码请求 ID
    - **expires_in**: 过期时间(秒)
    """
    # 生成随机字符串
    chars = CaptchaConfig.CHARSET
    captcha_text = ''.join(random.choices(chars, k=length))
    
    # 存储到内存
    request_id = f"captcha_{int(__import__('time').time() * 1000)}"
    _active_captchas[request_id] = CaptchaResult(captcha_text)
    
    # 如果请求使用 Cookie，设置 HttpOnly Cookie
    if use_cookie:
        response.set_cookie(
            key="llm_captcha_text",
            value=captcha_text,
            max_age=CaptchaConfig.EXPIRE_SECONDS,
            httponly=False,  # 允许前端读取
            secure=False,    # HTTP/HTTPS 都适用
            samesite="Lax",  # 防止 CSRF
            domain=".imageteach.tech",
            path="/"
        )
    
    return {
        "request_id": request_id,
        "expires_in": CaptchaConfig.EXPIRE_SECONDS,
        "cookie_set": use_cookie
    }


@router.post("/validate", summary="验证用户输入的验证码")
async def validate_captcha(
    request: Request,
    user_input: str,
    request_id: Optional[str] = None,
):
    """
    验证用户输入的验证码
    
    - **user_input**: 用户输入的验证码文本
    - **request_id**: 可选，指定验证哪个验证码
    
    返回:
    - **is_valid**: 是否验证成功
    - **submitted**: 用户提交的文本
    """
    if not user_input:
        raise HTTPException(status_code=400, detail="缺少验证码输入")
    
    # 从请求头或 Cookie 获取 request_id
    rid = request_id or request.headers.get("X-Captcha-Request-Id")
    
    if not rid or rid not in _active_captchas:
        raise HTTPException(
            status_code=404,
            detail="未找到验证码，请先生成"
        )
    
    captcha_result = _active_captchas[rid]
    
    # 检查是否过期
    if captcha_result.is_expired:
        del _active_captchas[rid]
        raise HTTPException(
            status_code=410,  # Gone
            detail="验证码已过期，请重新生成"
        )
    
    # 不区分大小写比较
    is_valid = captcha_result.text.lower() == user_input.strip().lower()
    
    # 验证成功后删除该验证码 (一次性使用)
    if is_valid:
        del _active_captchas[rid]
    
    return {
        "is_valid": is_valid,
        "submitted": user_input,
        "matched": is_valid
    }


@router.delete("/clear", summary="清除验证码 Cookie")
async def clear_captcha_cookie(response: Response):
    """清除所有验证码相关的 Cookie"""
    response.delete_cookie(key="llm_captcha_text")
    response.delete_cookie(key="llm_captcha_result")
    return {"message": "验证码 Cookie 已清除"}


# ==================== Cookie 策略工具类 ====================

class CookiePolicy:
    """
    现代 Cookie 策略管理器
    
    支持的属性:
    - **HttpOnly**: 防止 XSS 读取 Cookie
    - **Secure**: 仅 HTTPS 传输 (当前设为 False 以支持 HTTP)
    - **SameSite**: Lax/Strict/None (控制跨域行为)
    - **Domain**: 指定域名 (.imageteach.tech)
    - **Path**: 路径 (/)
    - **Max-Age**: 过期时间 (秒)
    """
    
    # 默认 Cookie 配置
    DEFAULT_CONFIG = {
        "llm_session": {
            "max_age": 86400 * 7,      # 7天
            "httponly": True,           # 防止 XSS
            "secure": False,            # HTTP/HTTPS 都适用
            "samesite": "Lax",          # 防止 CSRF
            "domain": ".imageteach.tech",
            "path": "/"
        },
        "llm_captcha": {
            "max_age": 3600,            # 1小时
            "httponly": False,
            "secure": False,
            "samesite": "Lax",
            "domain": ".imageteach.tech",
            "path": "/"
        },
        "llm_preferences": {
            "max_age": 86400 * 30,      # 30天
            "httponly": False,
            "secure": True,
            "samesite": "None",         # 跨域场景
            "domain": ".imageteach.tech",
            "path": "/"
        }
    }
    
    @classmethod
    def set_cookie(cls, response: Response, name: str, value: str, override: dict = None):
        """
        设置 Cookie
        
        Args:
            response: FastAPI Response 对象
            name: Cookie 名称
            value: Cookie 值
            override: 可选的覆盖配置
        """
        config = cls.DEFAULT_CONFIG.get(name, {})
        if override:
            config.update(override)
        
        response.set_cookie(
            key=name,
            value=value,
            max_age=config.get("max_age", 3600),
            httponly=config.get("httponly", False),
            secure=config.get("secure", False),
            samesite=config.get("samesite", "Lax"),
            domain=config.get("domain", ".imageteach.tech"),
            path=config.get("path", "/")
        )
    
    @classmethod
    def get_cookie_config(cls, name: str) -> dict:
        """获取 Cookie 配置"""
        return cls.DEFAULT_CONFIG.get(name, {})


# ==================== 前端可用的 API ====================

@router.get("/cookies/list", summary="列出所有验证码相关 Cookie")
async def list_captcha_cookies(request: Request):
    """返回当前请求携带的所有验证码相关 Cookie"""
    cookies = {
        k: v for k, v in request.cookies.items() 
        if k.startswith("llm_")
    }
    return {
        "cookies": cookies,
        "has_captcha": "llm_captcha_text" in cookies,
        "has_session": "llm_session" in cookies
    }


@router.post("/cookies/set", summary="设置测试 Cookie")
async def set_test_cookie(
    response: Response,
    cookie_name: str = "llm_test",
    cookie_value: str = "test_value"
):
    """设置一个测试用的 Cookie"""
    CookiePolicy.set_cookie(response, cookie_name, cookie_value)
    return {
        "message": f"Cookie '{cookie_name}' 已设置",
        "name": cookie_name,
        "value": cookie_value
    }