dify/api/controllers/console/auth/forgot_password.py

import base64
import logging
import secrets

from flask_restful import Resource, reqparse

from controllers.console import api
from controllers.console.auth.error import (
    InvalidEmailError,
    InvalidTokenError,
    PasswordMismatchError,
    PasswordResetRateLimitExceededError,
)
from controllers.console.setup import setup_required
from extensions.ext_database import db
from libs.helper import email as email_validate
from libs.password import hash_password, valid_password
from models.account import Account
from services.account_service import AccountService
from services.errors.account import RateLimitExceededError


class ForgotPasswordSendEmailApi(Resource):

    @setup_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument('email', type=str, 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:
            try:
                AccountService.send_reset_password_email(account=account)
            except RateLimitExceededError:
                logging.warning(f"Rate limit exceeded for email: {account.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"}


class ForgotPasswordCheckApi(Resource):

    @setup_required
    def post(self):
        parser = reqparse.RequestParser()
        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)

        if reset_data is None:
            return {'is_valid': False, 'email': None}
        return {'is_valid': True, 'email': reset_data.get('email')}


class ForgotPasswordResetApi(Resource):

    @setup_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument('token', type=str, required=True, nullable=False, location='json')
        parser.add_argument('new_password', type=valid_password, required=True, nullable=False, location='json')
        parser.add_argument('password_confirm', type=valid_password, required=True, nullable=False, location='json')
        args = parser.parse_args()

        new_password = args['new_password']
        password_confirm = args['password_confirm']

        if str(new_password).strip() != str(password_confirm).strip():
            raise PasswordMismatchError()

        token = args['token']
        reset_data = AccountService.get_reset_password_data(token)

        if reset_data is None:
            raise InvalidTokenError()

        AccountService.revoke_reset_password_token(token)

        salt = secrets.token_bytes(16)
        base64_salt = base64.b64encode(salt).decode()

        password_hashed = hash_password(new_password, salt)
        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()

        return {'result': 'success'}


api.add_resource(ForgotPasswordSendEmailApi, '/forgot-password')
api.add_resource(ForgotPasswordCheckApi, '/forgot-password/validity')
api.add_resource(ForgotPasswordResetApi, '/forgot-password/resets')