feat: add reset password api

This commit is contained in:
Joe
2024-08-29 16:51:30 +08:00
parent da684ebfaa
commit ede775cb6a
8 changed files with 75 additions and 50 deletions

View File

@@ -31,7 +31,7 @@ class PasswordResetRateLimitExceededError(BaseHTTPException):
code = 429
class EmailLoginCodeError(BaseHTTPException):
error_code = "email_login_code_error"
description = "Email login code is invalid or expired."
class EmailCodeError(BaseHTTPException):
error_code = "email_code_error"
description = "Email code is invalid or expired."
code = 400

View File

@@ -2,10 +2,14 @@ import base64
import logging
import secrets
from flask import request
from flask_restful import Resource, reqparse
from configs import dify_config
from constants.languages import languages
from controllers.console import api
from controllers.console.auth.error import (
EmailCodeError,
InvalidEmailError,
InvalidTokenError,
PasswordMismatchError,
@@ -13,7 +17,7 @@ from controllers.console.auth.error import (
)
from controllers.console.setup import setup_required
from extensions.ext_database import db
from libs.helper import email as email_validate
from libs.helper import email, get_remote_ip
from libs.password import hash_password, valid_password
from models.account import Account
from services.account_service import AccountService
@@ -24,42 +28,48 @@ class ForgotPasswordSendEmailApi(Resource):
@setup_required
def post(self):
parser = reqparse.RequestParser()
parser.add_argument("email", type=str, required=True, location="json")
parser.add_argument("email", type=email, required=True, location="json")
args = parser.parse_args()
email = args["email"]
if not email_validate(email):
raise InvalidEmailError()
account = Account.query.filter_by(email=email).first()
if account:
account = Account.query.filter_by(email=args["email"]).first()
token = None
if account is None:
if dify_config.ALLOW_REGISTER:
token = AccountService.send_reset_password_email(email=args["email"])
else:
raise InvalidEmailError()
elif account:
try:
AccountService.send_reset_password_email(account=account)
token = AccountService.send_reset_password_email(account=account, email=args["email"])
except RateLimitExceededError:
logging.warning(f"Rate limit exceeded for email: {account.email}")
logging.warning(f"Rate limit exceeded for email: {args["email"]}")
raise PasswordResetRateLimitExceededError()
else:
# Return success to avoid revealing email registration status
logging.warning(f"Attempt to reset password for unregistered email: {email}")
return {"result": "success"}
return {"result": "success", "data": token}
class ForgotPasswordCheckApi(Resource):
@setup_required
def post(self):
parser = reqparse.RequestParser()
parser.add_argument("email", type=str, required=True, location="json")
parser.add_argument("code", type=str, required=True, location="json")
parser.add_argument("token", type=str, required=True, nullable=False, location="json")
args = parser.parse_args()
token = args["token"]
reset_data = AccountService.get_reset_password_data(token)
user_email = args["email"]
if reset_data is None:
return {"is_valid": False, "email": None}
return {"is_valid": True, "email": reset_data.get("email")}
token_data = AccountService.get_reset_password_data(args["token"])
if token_data is None:
raise InvalidTokenError()
if user_email != token_data.get("email"):
raise InvalidEmailError()
if args["code"] != token_data.get("code"):
raise EmailCodeError()
return {"is_valid": True, "email": token_data.get("email")}
class ForgotPasswordResetApi(Resource):
@@ -92,11 +102,21 @@ class ForgotPasswordResetApi(Resource):
base64_password_hashed = base64.b64encode(password_hashed).decode()
account = Account.query.filter_by(email=reset_data.get("email")).first()
account.password = base64_password_hashed
account.password_salt = base64_salt
db.session.commit()
if account:
account.password = base64_password_hashed
account.password_salt = base64_salt
db.session.commit()
else:
account = AccountService.create_user_through_env(
email=reset_data.get("email"),
name=reset_data.get("email"),
password=password_confirm,
interface_language=languages[0],
)
return {"result": "success"}
token = AccountService.login(account, ip_address=get_remote_ip(request))
return {"result": "success", "data": token}
api.add_resource(ForgotPasswordSendEmailApi, "/forgot-password")

View File

@@ -1,7 +1,7 @@
from typing import cast
import flask_login
from flask import request
from flask import redirect, request
from flask_restful import Resource, reqparse
import services
@@ -9,7 +9,7 @@ from configs import dify_config
from constants.languages import languages
from controllers.console import api
from controllers.console.auth.error import (
EmailLoginCodeError,
EmailCodeError,
InvalidEmailError,
InvalidTokenError,
)
@@ -134,13 +134,14 @@ class EmailCodeLoginSendEmailApi(Resource):
@setup_required
def post(self):
parser = reqparse.RequestParser()
parser.add_argument("email", type=str, required=True, location="json")
parser.add_argument("email", type=email, required=True, location="json")
args = parser.parse_args()
account = AccountService.get_user_through_email(args["email"])
if account is None:
if dify_config.ALLOW_REGISTER:
token = AccountService.send_email_code_login_email(email=args["email"])
token = AccountService.send_reset_password_email(email=args["email"])
return redirect(f"{dify_config.CONSOLE_WEB_URL}/reset-password?token={token}")
else:
raise InvalidEmailError()
else:
@@ -168,7 +169,7 @@ class EmailCodeLoginApi(Resource):
raise InvalidEmailError()
if token_data["code"] != args["code"]:
raise EmailLoginCodeError()
raise EmailCodeError()
AccountService.revoke_email_code_login_token(args["token"])
account = AccountService.get_user_through_email(user_email)
@@ -186,4 +187,4 @@ api.add_resource(LoginApi, "/login")
api.add_resource(LogoutApi, "/logout")
api.add_resource(EmailCodeLoginSendEmailApi, "/email-code-login")
api.add_resource(EmailCodeLoginApi, "/email-code-login/validity")
api.add_resource(ResetPasswordApi, "/reset-password")
api.add_resource(ResetPasswordSendEmailApi, "/reset-password")

View File

@@ -190,8 +190,11 @@ def compact_generate_response(response: Union[dict, RateLimitGenerator]) -> Resp
class TokenManager:
@classmethod
def generate_token(
cls, token_type: str, account: Optional[Account] = None, email: Optional[str] = None,
additional_data: dict = None
cls,
token_type: str,
account: Optional[Account] = None,
email: Optional[str] = None,
additional_data: dict = None,
) -> str:
if account is None and email is None:
raise ValueError("Account or email must be provided")

View File

@@ -250,19 +250,22 @@ class AccountService:
@classmethod
def send_reset_password_email(cls, account: Optional[Account] = None, email: Optional[str] = None):
account_email = account.email if account else email
account_language = account.interface_language if account else languages[0]
if cls.reset_password_rate_limiter.is_rate_limited(account.email):
raise RateLimitExceededError(f"Rate limit exceeded for email: {account.email}. Please try again later.")
raise RateLimitExceededError(f"Rate limit exceeded for email: {account_email}. Please try again later.")
code = "".join([str(random.randint(0, 9)) for _ in range(6)])
token = TokenManager.generate_token(
account=account, email=email, token_type="reset_password", additional_data={"code": code}
)
send_reset_password_mail_task.delay(
language=account.interface_language if account else languages[0],
to=account.email if account else email,
language=account_language,
to=account_email,
code=code,
)
cls.reset_password_rate_limiter.increment_rate_limit(account.email if account else email)
cls.reset_password_rate_limiter.increment_rate_limit(account_email)
return token
@classmethod

View File

@@ -5,17 +5,16 @@ import click
from celery import shared_task
from flask import render_template
from configs import dify_config
from extensions.ext_mail import mail
@shared_task(queue="mail")
def send_reset_password_mail_task(language: str, to: str, token: str):
def send_reset_password_mail_task(language: str, to: str, code: str):
"""
Async Send reset password mail
:param language: Language in which the email should be sent (e.g., 'en', 'zh')
:param to: Recipient email address
:param token: Reset password token to be included in the email
:param code: Reset password code
"""
if not mail.is_inited():
return
@@ -25,12 +24,11 @@ def send_reset_password_mail_task(language: str, to: str, token: str):
# send reset password mail using different languages
try:
url = f"{dify_config.CONSOLE_WEB_URL}/forgot-password?token={token}"
if language == "zh-Hans":
html_content = render_template("reset_password_mail_template_zh-CN.html", to=to, url=url)
html_content = render_template("reset_password_mail_template_zh-CN.html", to=to, code=code)
mail.send(to=to, subject="重置您的 Dify 密码", html=html_content)
else:
html_content = render_template("reset_password_mail_template_en-US.html", to=to, url=url)
html_content = render_template("reset_password_mail_template_en-US.html", to=to, code=code)
mail.send(to=to, subject="Reset Your Dify Password", html=html_content)
end_at = time.perf_counter()

View File

@@ -63,12 +63,12 @@
<!-- Optional: Add a logo or a header image here -->
<img src="https://cloud.dify.ai/logo/logo-site.png" alt="Dify Logo" />
</div>
<p class="title">Reset your Dify password</p>
<p class="title">Reset your Dify password</p>
<p class="description">Copy and paste this code, this code will only be valid for the next 5 minutes.</p>
<div class="code-content">
<span class="code">{{code}}</span>
</div>
<p class="tips">If you didn't request a reset, don't worry. You can safely ignore this email.</p>
<p class="tips">If you didn't request a reset, don't worry. You can safely ignore this email.</p>
</div>
</body>
</html>

View File

@@ -63,12 +63,12 @@
<!-- Optional: Add a logo or a header image here -->
<img src="https://cloud.dify.ai/logo/logo-site.png" alt="Dify Logo" />
</div>
<p class="title">重置您的Dify密码</p>
<p class="title">重置您的Dify 账户密码</p>
<p class="description">复制并粘贴此验证码,注意验证码仅在接下来的 5 分钟内有效。</p>
<div class="code-content">
<span class="code">{{code}}</span>
</div>
<p class="tips">如果您没有请求重置密码,请不要担心。您可以安全地忽略此电子邮件。</p>
<p class="tips">如果您没有请求重置,请不要担心。您可以安全地忽略此电子邮件。</p>
</div>
</body>
</html>