From ede775cb6a24f6d2ec004a157353bc7447557eaf Mon Sep 17 00:00:00 2001
From: Joe <1264204425@qq.com>
Date: Thu, 29 Aug 2024 16:51:30 +0800
Subject: [PATCH] feat: add reset password api
---
api/controllers/console/auth/error.py | 6 +-
.../console/auth/forgot_password.py | 70 ++++++++++++-------
api/controllers/console/auth/login.py | 13 ++--
api/libs/helper.py | 7 +-
api/services/account_service.py | 11 +--
api/tasks/mail_reset_password_task.py | 10 ++-
.../reset_password_mail_template_en-US.html | 4 +-
.../reset_password_mail_template_zh-CN.html | 4 +-
8 files changed, 75 insertions(+), 50 deletions(-)
diff --git a/api/controllers/console/auth/error.py b/api/controllers/console/auth/error.py
index 1c85035d25..98b4e96beb 100644
--- a/api/controllers/console/auth/error.py
+++ b/api/controllers/console/auth/error.py
@@ -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
diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py
index 0b01a4906a..8df148538a 100644
--- a/api/controllers/console/auth/forgot_password.py
+++ b/api/controllers/console/auth/forgot_password.py
@@ -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")
diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py
index 6e439863b8..478d7b9c8a 100644
--- a/api/controllers/console/auth/login.py
+++ b/api/controllers/console/auth/login.py
@@ -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")
diff --git a/api/libs/helper.py b/api/libs/helper.py
index b39e83dc03..7e3c269e3f 100644
--- a/api/libs/helper.py
+++ b/api/libs/helper.py
@@ -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")
diff --git a/api/services/account_service.py b/api/services/account_service.py
index e2f30d8ed3..b1912a6e70 100644
--- a/api/services/account_service.py
+++ b/api/services/account_service.py
@@ -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
diff --git a/api/tasks/mail_reset_password_task.py b/api/tasks/mail_reset_password_task.py
index cbb78976ca..7a0d2c877b 100644
--- a/api/tasks/mail_reset_password_task.py
+++ b/api/tasks/mail_reset_password_task.py
@@ -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()
diff --git a/api/templates/reset_password_mail_template_en-US.html b/api/templates/reset_password_mail_template_en-US.html
index fa73144ad4..da8a383ea1 100644
--- a/api/templates/reset_password_mail_template_en-US.html
+++ b/api/templates/reset_password_mail_template_en-US.html
@@ -63,12 +63,12 @@
- Reset your Dify password
+ Reset your Dify password
Copy and paste this code, this code will only be valid for the next 5 minutes.
{{code}}
- If you didn't request a reset, don't worry. You can safely ignore this email.
+ If you didn't request a reset, don't worry. You can safely ignore this email.