Python/agent/application/states/auth.py

# -*- coding: utf-8 -*-
"""
认证状态
"""
from asyncio import sleep, create_task
from email.header import Header
from email.mime.text import MIMEText
import re
from typing import Any, AsyncGenerator, Coroutine, cast
from aiosmtplib import SMTP, SMTPException
import reflex as rx

from application.states.database import DatabaseState

# 验证码HTML模板
CAPTCHA_HTML_TEMPLATE = """
<p>验证码：<strong>{captcha}</strong></p>
"""


class AuthState(rx.State):
    """
    认证状态
    """

    # 用户唯一标识（通过本地存储同步）
    user_id: str = rx.LocalStorage("user_id", sync=True)

    # 邮箱
    email: str = ""
    # 验证码
    captcha: str = ""
    # 协议同意状态
    is_policies_agreed: bool = False
    # 提示信息
    notification: str = ""

    # 验证码发送倒计时
    captcha_sending_countdown: int = 0

    # 登录中状态
    is_logging_in: bool = False

    # 激活的导航按钮
    activated_nav_button: str = "conversation"

    @rx.event
    async def check(self) -> None:
        """
        检查
        """

        self.user_id = ""

    @rx.event
    def set_email(self, email: str) -> None:
        """
        设置邮箱
        :param email: 邮箱
        :return: None
        """
        self.email = email.strip()

    @rx.event
    def set_captcha(self, captcha: str) -> None:
        """
        设置验证码
        :param captcha: 验证码
        :return: None
        """
        self.captcha = captcha.strip()

    @rx.event
    def toggle_policies_agreed(self) -> None:
        """
        切换协议同意状态
        :return: None
        """
        self.is_policies_agreed = not self.is_policies_agreed

    def validate_email(self) -> bool:
        """
        验证邮箱
        :return: 验证是否通过，True 表示校验通过，False 表示校验不通过
        """
        self.notification = ""
        # 非空校验
        if not self.email:
            self.notification = "请输入邮箱"
            return False
        # 格式校验
        if not re.fullmatch(
            r"^[a-zA-Z0-9]+([_\-+.][a-zA-Z0-9_\-+]+)*@[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*(\.[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*)*\.[a-zA-Z0-9]{2,}$",
            self.email,
        ):
            self.notification = "邮箱格式不正确，请重新输入"
            return False
        return True

    def validate_captcha(self) -> bool:
        """
        验证验证码
        :return: 验证是否通过，True 表示校验通过，False 表示校验不通过
        """
        self.notification = ""
        # 非空校验
        if not self.captcha:
            self.notification = "请输入验证码"
            return False
        # 格式校验
        if not re.fullmatch(r"^\d{6}$", self.captcha):
            self.notification = "验证码格式不正确，请重新输入"
            return False
        return True

    @rx.event(background=True)
    async def resend_captcha_countdown(self) -> None:
        """
        重新发送验证码倒计时
        :return: None
        """

        async with self:
            self.captcha_sending_countdown = 59

        while self.captcha_sending_countdown > 0:
            await sleep(1)
            async with self:
                self.captcha_sending_countdown -= 1

    @rx.event(background=True)
    async def send_captcha_(self) -> None:
        """
        发送验证码
        :return: None
        """
        async with self:
            self.notification = "验证码已发送"
            # 创建验证码记录并赋值
            database_state = await self.get_state(DatabaseState)
            captcha = await database_state.create_captchas_record(email=self.email)

        # 构建邮件
        message = MIMEText(
            CAPTCHA_HTML_TEMPLATE.format(captcha=captcha),
            "html",
            "utf-8",
        )
        message["Subject"] = Header("验证码", "utf-8").encode()
        message["From"] = "office@liubiren.cloud"
        message["To"] = self.email
        try:
            async with SMTP(
                hostname="smtp.feishu.cn",
                port=465,
                username="office@liubiren.cloud",
                password="cx1dM5tBcbyWftpS",
                use_tls=True,
                local_hostname="localhost",
            ) as smtp:
                await smtp.send_message(message)

        except SMTPException:
            async with self:
                self.notification = f"发送失败，请稍后重试"
                self.captcha_sending_countdown = 0

    @rx.event
    async def send_captcha(self) -> AsyncGenerator:
        """
        发送验证码
        :return: None
        """
        self.notification = ""

        # 验证邮箱
        if not self.validate_email():
            return

        # 将重新发送验证码倒计时和发送验证码事件添加至后台任务队列
        yield type(self).resend_captcha_countdown
        yield type(self).send_captcha_

    @rx.event
    async def login(self) -> None:
        """
        登录
        :return: None
        """
        self.is_logging_in = True

        # 验证邮箱
        if not self.validate_email():
            return

        # 验证验证码
        if not self.validate_captcha():
            return

        # 创建验证码记录并赋值验证码
        database_state = await self.get_state(DatabaseState)

        if not await database_state.update_captcha(
            email=self.email, captcha=self.captcha
        ):
            self.notification = "验证码错误"
            return

        # 创建用户记录
        user_id = await database_state.create_users_record(email=self.email)
        self.user_id = user_id
        self.notification = "登录成功"
        self.is_logging_in = False

    @rx.event
    def switch_activated_nav_button(self, button: str):
        """
        将指定导航按钮设置为激活的导航按钮
        """
        self.activated_nav_button = button