zfsafe/web/common/utils.py

from PIL import Image, ImageDraw, ImageFont,ImageFilter
import random
import string
import io
from functools import wraps
from quart import session, redirect, url_for, flash,current_app

def generate_captcha():
    characters = string.ascii_uppercase + string.digits
    captcha_text = ''.join(random.choices(characters, k=6))

    font = ImageFont.truetype("arial.ttf", 36)
    image = Image.new('RGB', (200, 60), color=(255, 255, 255))
    draw = ImageDraw.Draw(image)

    for i in range(6):
        draw.text((10 + i * 30, 10), captcha_text[i], font=font,
                  fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    # 模糊化处理
    image = image.filter(ImageFilter.BLUR)

    # 将图片保存到BytesIO流中
    buffer = io.BytesIO()
    image.save(buffer, 'jpeg')
    buffer.seek(0)

    # captcha_path = os.path.join('static', 'captcha', f'{captcha_text}.png')
    # image.save(captcha_path)
    # return captcha_text, f'static/captcha/{captcha_text}.png'
    return captcha_text, buffer


def verify_captcha(user_input, actual_captcha):
    return user_input == actual_captcha


def login_required(f):
    @wraps(f)
    async def decorated_function(*args, **kwargs):
        username = session.get('user')
        token = session.get('token')
        if not username or not token:
            await flash('未登录，请重新登录', 'error')
            return redirect(url_for('main.login'))
        #从redis取最新的token
        redis_key = f"user_token:{username}"
        server_token = await current_app.redis.get(redis_key)
        if server_token is None or server_token != token:
            # 如果不匹配，说明此账号已在其他地方登录
            session.clear()
            await flash('您的账号已在其他设备登录，请重新登录', 'error')
            return redirect(url_for('main.login'))

        return await f(*args, **kwargs)
    return decorated_function