97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
import json
|
|
import uuid
|
|
import requests
|
|
from enum import Enum
|
|
from captcha.image import ImageCaptcha
|
|
from Utils.EmailUtils import send_email
|
|
from Utils.RandomUtils import get_random_num_code, get_random_letter_and_num_code
|
|
from Utils.RedisUtils import RedisPool
|
|
|
|
|
|
class EmailVerifyType(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 EmailCodeVerify:
|
|
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
|