新增重置密码功能

app/core/executor.py

@@ -666,7 +666,7 @@ async def notice(env: list, plan: PityTestPlan, project: Project, report_dict: d
                 for m in msg_types:
                     if int(m) == NoticeType.EMAIL:
                         render_html = Email.render_html(plan_name=plan.name, **report_dict[e])
-                        Email.send_msg(
+                        await Email.send_msg(
                             f"【{report_dict[e].get('env')}】测试计划【{plan.name}】执行完毕（{report_dict[e].get('plan_result')}）",
                             render_html, None, *[r.get("email") for r in users])
                     if int(m) == NoticeType.DINGDING:

app/core/msg/mail.py

@@ -3,6 +3,8 @@
 from email.mime.text import MIMEText
 from email.utils import make_msgid
 
+import aioify
+from awaits.awaitable import awaitable
 from jinja2.environment import Template
 
 from app.core.configuration import SystemConfiguration
@@ -24,6 +26,7 @@ class Email(Notification):
     #     client.send(receiver, subject=subject, contents=content, attachments=attachment)
 
     @staticmethod
+    @awaitable
     def send_msg(subject, content, attachment=None, *receiver):
         configuration = SystemConfiguration.get_config()
         data = configuration.get("email")

app/crud/auth/UserDao.py

@@ -1,4 +1,3 @@
-import asyncio
 import random
 import time
 from datetime import datetime
@@ -9,7 +8,7 @@
 from app.crud import Mapper
 from app.middleware.Jwt import UserToken
 from app.middleware.RedisManager import RedisHelper
-from app.models import async_session, DatabaseHelper
+from app.models import async_session
 from app.models.user import User
 from app.schema.user import UserUpdateForm
 from app.utils.logger import Log
@@ -149,7 +148,7 @@ async def login(username, password):
                 async with session.begin():
                     # 查询用户名/密码匹配且没有被删除的用户
                     query = await session.execute(
-                        select(User).where(User.username == username, User.password == pwd,
+                        select(User).where(or_(User.username == username, User.email == username), User.password == pwd,
                                            User.deleted_at == 0))
                     user = query.scalars().first()
                     if user is None:
@@ -195,3 +194,22 @@ async def list_user_touch(*user):
         except Exception as e:
             UserDao.log.error(f"获取用户联系方式失败: {str(e)}")
             raise Exception(f"获取用户联系方式失败: {e}")
+
+    @staticmethod
+    async def reset_password(email: str, password: str):
+        pwd = UserToken.add_salt(password)
+        try:
+            async with async_session() as session:
+                async with session.begin():
+                    sql = update(User).where(User.email == email).values(password=pwd)
+                    await session.execute(sql)
+        except Exception as e:
+            UserDao.log.error(f"重置用户: {email}密码失败: {str(e)}")
+            raise Exception(f"重置{email}密码失败")
+
+    @staticmethod
+    async def query_user_by_email(email: str):
+        async with async_session() as session:
+            sql = select(User).where(User.email == email, User.is_valid == True)
+            query = await session.execute(sql)
+            return query.scalars().first()

app/routers/auth/user.py

@@ -1,14 +1,17 @@
+import asyncio
+
 import requests
 from fastapi import APIRouter, Depends
 from starlette import status
 
+from app.core.msg.mail import Email
 from app.crud.auth.UserDao import UserDao
 from app.excpetions.RequestException import AuthException
 from app.handler.fatcory import PityResponse
 from app.middleware.Jwt import UserToken
 from app.routers import Permission, FORBIDDEN
-from app.routers.auth.user_schema import UserDto, UserForm
-from app.schema.user import UserUpdateForm
+from app.schema.user import UserUpdateForm, UserForm, UserDto, ResetPwdForm
+from app.utils.des import Des
 from config import Config
 
 router = APIRouter(prefix="/auth")
@@ -106,3 +109,34 @@ async def delete_user(id: int, user=Depends(Permission(Config.ADMIN))):
         return PityResponse.success(user)
     except Exception as e:
         return PityResponse.failed(e)
+
+
+@router.post("/reset", summary="重置用户密码")
+async def reset_user(form: ResetPwdForm):
+    email = Des.des_decrypt(form.token)
+    await UserDao.reset_password(email, form.password)
+    return PityResponse.success()
+
+
+@router.get("/reset/generate/{email}", summary="生成重置密码链接")
+async def generate_reset_url(email: str):
+    try:
+        user = await UserDao.query_user_by_email(email)
+        if user is not None:
+            # 说明邮件存在，发送邮件
+            em = Des.des_encrypt(email)
+            link = f"""https://pity.fun/#/user/resetPassword?token={em}"""
+            render_html = Email.render_html(Config.PASSWORD_HTML_PATH, link=link, name=user.name)
+            asyncio.create_task(Email.send_msg("重置你的pity密码", render_html, None, email))
+        return PityResponse.success(None)
+    except Exception as e:
+        return PityResponse.failed(str(e))
+
+
+@router.get("/reset/check/{token}", summary="检测生成的链接是否正确")
+async def check_reset_url(token: str):
+    try:
+        email = Des.des_decrypt(token)
+        return PityResponse.success(email)
+    except:
+        return PityResponse.failed("重置链接不存在, 请不要无脑尝试")

app/schema/user.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel, validator
 
 # 都可以为空，为空则不进行更改
+from app.excpetions.ParamsException import ParamsError
 from app.schema.base import PityModel
 
 
@@ -15,3 +16,38 @@ class UserUpdateForm(BaseModel):
     @validator('id')
     def id_not_empty(cls, v):
         return PityModel.not_empty(v)
+
+
+class UserDto(BaseModel):
+    name: str
+    password: str
+    username: str
+    email: str
+
+    @validator('name', 'password', 'username', 'email')
+    def field_not_empty(cls, v):
+        if isinstance(v, str) and len(v.strip()) == 0:
+            raise ParamsError("不能为空")
+        return v
+
+
+class UserForm(BaseModel):
+    username: str
+    password: str
+
+    @validator('password', 'username')
+    def name_not_empty(cls, v):
+        if isinstance(v, str) and len(v.strip()) == 0:
+            raise ParamsError("不能为空")
+        return v
+
+
+class ResetPwdForm(BaseModel):
+    password: str
+    token: str
+
+    @validator('token', 'password')
+    def name_not_empty(cls, v):
+        if isinstance(v, str) and len(v.strip()) == 0:
+            raise ParamsError("不能为空")
+        return v

app/utils/des.py

@@ -0,0 +1,38 @@
+import binascii
+
+from pyDes import des, CBC, PAD_PKCS5
+
+# 秘钥
+KEY = 'pityspwd'
+
+
+class Des(object):
+
+    @staticmethod
+    def des_encrypt(s):
+        """
+        DES 加密
+        :param s: 原始字符串
+        :return: 加密后字符串，16进制
+        """
+        secret_key = KEY  # 密码
+        iv = secret_key  # 偏移
+        # secret_key:加密密钥，CBC:加密模式，iv:偏移, padmode:填充
+        des_obj = des(secret_key, CBC, iv, pad=None, padmode=PAD_PKCS5)
+        # 返回为字节
+        secret_bytes = des_obj.encrypt(s, padmode=PAD_PKCS5)
+        # 返回为16进制
+        return binascii.b2a_hex(secret_bytes).decode()
+
+    @staticmethod
+    def des_decrypt(s):
+        """
+        DES 解密
+        :param s: 加密后的字符串，16进制
+        :return:  解密后的字符串
+        """
+        secret_key = KEY
+        iv = secret_key
+        des_obj = des(secret_key, CBC, iv, pad=None, padmode=PAD_PKCS5)
+        decrypt_str = des_obj.decrypt(binascii.a2b_hex(s), padmode=PAD_PKCS5)
+        return decrypt_str.decode()

config.py

@@ -61,6 +61,9 @@ class BaseConfig(BaseSettings):
     # 测试报告路径
     REPORT_PATH = os.path.join(ROOT, "templates", "report.html")
 
+    # 重置密码路径
+    PASSWORD_HTML_PATH = os.path.join(ROOT, "templates", "reset_password.html")
+
     # APP 路径
     APP_PATH = os.path.join(ROOT, "app")
 

requirements.txt

@@ -33,4 +33,5 @@ jsonpath~=0.82
 py-mock
 starlette_context
 gunicorn
-supervisor
+supervisor
+pyDes~=2.0.1