wd-rating/utils/verify_code_utils.py

import json
import uuid
import requests
from enum import Enum
from captcha.image import ImageCaptcha
from utils.email_utils import send_email
from utils.random_utils import get_random_num_code, get_random_letter_and_num_code
from utils.redis_utils import RedisPool


class EmailVerifyType(Enum):
    login = 1
    register = 2
    reset_password = 3
    change = 4


class PhoneVerifyType(Enum):
    login = 1
    register = 2
    reset_password = 3
    change = 4


class ImageCaptchaVerify:
    def __init__(self, redis_pool: RedisPool):
        self.redis_pool = redis_pool

    def make_captcha_image(self, code=None, expire_time_s: int = 120):
        captcha_id = str(uuid.uuid4())
        image = ImageCaptcha()
        if not code:
            code = get_random_letter_and_num_code(4)
        image_bytes = image.generate(code).getvalue()
        self.redis_pool.conn.set(captcha_id, code, expire_time_s)
        return captcha_id, image_bytes, code

    def check_code(self, captcha_id, code: str) -> bytes:
        captcha_code = self.redis_pool.conn.get(captcha_id)
        self.redis_pool.conn.delete(captcha_id)
        if captcha_code and captcha_code.lower() != code.lower():
            return False
        return True


class EmailVerifyCode:
    def __init__(self, redis_pool: RedisPool, email_api: str, email_sender: str, id_sign="EmailVerifyCodeId"):
        self.redis_pool = redis_pool
        self.email_api = email_api
        self.email_sender = email_sender
        self.id_sign = id_sign

    def get_email_id(self, email, verify_type: EmailVerifyType = EmailVerifyType.login):
        """
        :param email:
        :param verify_type: 不同业务验证类型对应不同的id
        :return:str
        """
        return f"{email}_{self.id_sign}_{verify_type.name}"

    def make_code(self, email, expire_time_s=60 * 5, code_len=4, verify_type: EmailVerifyType = EmailVerifyType.login):
        client = self.redis_pool.get_redis_client()
        email_id = self.get_email_id(email, verify_type)
        code = get_random_num_code(code_len)
        client.set(email_id, code)
        if expire_time_s:
            client.expire(email_id, expire_time_s)
        return code

    def send_code(self, email, verify_type=EmailVerifyType.register):
        code = self.make_code(email, verify_type=verify_type)
        title = ""
        if verify_type == EmailVerifyType.login:
            title = "登录验证码"
        elif verify_type == EmailVerifyType.register:
            title = "注册验证码"
        elif verify_type == EmailVerifyType.change:
            title = "修改信息验证码"
        elif verify_type == EmailVerifyType.reset_password:
            title = "修改密码验证码"
        send_email(title, email, code)

    def check_code(self, email, code, verify_type: EmailVerifyType = EmailVerifyType.login):
        email_id = self.get_email_id(email, verify_type)
        client = self.redis_pool.get_redis_client()
        email_code = client.get(email_id)
        return code and code == email_code

    def remove_code(self, email, verify_type: EmailVerifyType = EmailVerifyType.login):
        client = self.redis_pool.get_redis_client()
        email_id = self.get_email_id(email, verify_type=verify_type)
        client.delete(email_id)

    def send_email(self, title: str, email: str, msg_body: str):
        headers = {"Content-Type": "application/json;charset=UTF-8"}
        data = {
            "title": title,
            "sender": self.sender,
            "recipients": [email],
            "msg_body": msg_body
        }
        res = requests.post(url=self.email_api + '/send_mail', headers=headers, data=json.dumps(data), timeout=5)
        return res.status_code == 200


class PhoneVerifyCode:
    def __init__(self, redis_pool: RedisPool, phone_msg_send, id_sign="PhoneVerifyCodeId"):
        self.redis_pool = redis_pool
        self.phone_msg_send = phone_msg_send
        self.id_sign = id_sign

    def get_phone_id(self, phone, verify_type: PhoneVerifyType = PhoneVerifyType.login):
        """
        :param phone:
        :param verify_type: 不同业务验证类型对应不同的id
        :return:str
        """
        return f"{phone}_{self.id_sign}_{verify_type.name}"

    def make_code(self, phone, expire_time_s=60 * 5, verify_type: PhoneVerifyType = PhoneVerifyType.login):
        client = self.redis_pool.get_redis_client()
        phone_id = self.get_phone_id(phone, verify_type)
        code = get_random_num_code(4)
        client.set(phone_id, code)
        if expire_time_s:
            client.expire(phone_id, expire_time_s)
        return code

    def send_code(self, phone, verify_type: PhoneVerifyType):
        code = self.make_code(phone, verify_type=verify_type)
        self.phone_msg_send.send_phone_code(phone, code)

    def check_code(self, phone, code, verify_type: PhoneVerifyType = PhoneVerifyType.login):
        phone_id = self.get_phone_id(phone, verify_type=verify_type)
        client = self.redis_pool.get_redis_client()
        phone_code = client.get(phone_id)
        return code and code == phone_code