73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
from io import BytesIO
|
||
from captcha.image import ImageCaptcha
|
||
from uuid import uuid4
|
||
from utils.redis_client import get_redis, set_token_in_redis
|
||
import redis.asyncio as redis
|
||
import random
|
||
import string
|
||
from fastapi import HTTPException
|
||
|
||
|
||
def generate_captcha(length: int = 4) -> tuple[str, bytes]:
|
||
"""生成验证码图片和验证码文本"""
|
||
# 生成随机验证码
|
||
characters = string.digits + string.ascii_uppercase
|
||
captcha_text = ''.join(random.choices(characters, k=length))
|
||
|
||
try:
|
||
# 生成验证码图片
|
||
image = ImageCaptcha(width=150, height=50)
|
||
image_bytes = image.generate(captcha_text)
|
||
|
||
# 将图片转换为字节流
|
||
buffer = BytesIO()
|
||
buffer.write(image_bytes.read())
|
||
buffer.seek(0)
|
||
|
||
return captcha_text, buffer.getvalue()
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail="验证码生成失败")
|
||
|
||
|
||
async def save_captcha_to_redis(redis_conn: redis.Redis, captcha_id: str, captcha_text: str, expire_seconds: int = 300):
|
||
"""将验证码保存到Redis"""
|
||
await redis_conn.set(f"captcha:{captcha_id}", captcha_text, ex=expire_seconds)
|
||
|
||
|
||
async def get_captcha_from_redis(redis_conn: redis.Redis, captcha_id: str) -> str | None:
|
||
"""从Redis获取验证码"""
|
||
return await redis_conn.get(f"captcha:{captcha_id}")
|
||
|
||
|
||
async def delete_captcha_from_redis(redis_conn: redis.Redis, captcha_id: str):
|
||
"""从Redis删除验证码"""
|
||
await redis_conn.delete(f"captcha:{captcha_id}")
|
||
|
||
|
||
async def validate_captcha(redis_conn: redis.Redis, captcha_id: str, user_input: str) -> bool:
|
||
"""验证验证码是否正确"""
|
||
stored_captcha = await get_captcha_from_redis(redis_conn, captcha_id)
|
||
if not stored_captcha:
|
||
return False
|
||
|
||
# 忽略大小写比较
|
||
is_valid = stored_captcha.upper() == user_input.upper()
|
||
|
||
# 无论验证成功与否，都删除验证码（防止重复使用）
|
||
await delete_captcha_from_redis(redis_conn, captcha_id)
|
||
|
||
return is_valid
|
||
|
||
|
||
async def create_captcha(redis_conn: redis.Redis) -> tuple[str, bytes]:
|
||
"""创建验证码并保存到Redis"""
|
||
# 生成验证码
|
||
captcha_text, image_bytes = generate_captcha()
|
||
|
||
# 生成唯一ID
|
||
captcha_id = str(uuid4())
|
||
|
||
# 保存到Redis，有效期5分钟
|
||
await save_captcha_to_redis(redis_conn, captcha_id, captcha_text, 300)
|
||
|
||
return captcha_id, image_bytes |