Underground-Digital/Workflow-Engine

1import base642import secrets3 4from flask import request5from flask_restful import Resource, reqparse6 7from constants.languages import languages8from controllers.console import api9from controllers.console.auth.error import (10    EmailCodeError,11    InvalidEmailError,12    InvalidTokenError,13    PasswordMismatchError,14)15from controllers.console.error import EmailSendIpLimitError, NotAllowedRegister16from controllers.console.wraps import setup_required17from events.tenant_event import tenant_was_created18from extensions.ext_database import db19from libs.helper import email, extract_remote_ip20from libs.password import hash_password, valid_password21from models.account import Account22from services.account_service import AccountService, TenantService23from services.errors.workspace import WorkSpaceNotAllowedCreateError24from services.feature_service import FeatureService25 26 27class ForgotPasswordSendEmailApi(Resource):28    @setup_required29    def post(self):30        parser = reqparse.RequestParser()31        parser.add_argument("email", type=email, required=True, location="json")32        parser.add_argument("language", type=str, required=False, location="json")33        args = parser.parse_args()34 35        ip_address = extract_remote_ip(request)36        if AccountService.is_email_send_ip_limit(ip_address):37            raise EmailSendIpLimitError()38 39        if args["language"] is not None and args["language"] == "zh-Hans":40            language = "zh-Hans"41        else:42            language = "en-US"43 44        account = Account.query.filter_by(email=args["email"]).first()45        token = None46        if account is None:47            if FeatureService.get_system_features().is_allow_register:48                token = AccountService.send_reset_password_email(email=args["email"], language=language)49                return {"result": "fail", "data": token, "code": "account_not_found"}50            else:51                raise NotAllowedRegister()52        else:53            token = AccountService.send_reset_password_email(account=account, email=args["email"], language=language)54 55        return {"result": "success", "data": token}56 57 58class ForgotPasswordCheckApi(Resource):59    @setup_required60    def post(self):61        parser = reqparse.RequestParser()62        parser.add_argument("email", type=str, required=True, location="json")63        parser.add_argument("code", type=str, required=True, location="json")64        parser.add_argument("token", type=str, required=True, nullable=False, location="json")65        args = parser.parse_args()66 67        user_email = args["email"]68 69        token_data = AccountService.get_reset_password_data(args["token"])70        if token_data is None:71            raise InvalidTokenError()72 73        if user_email != token_data.get("email"):74            raise InvalidEmailError()75 76        if args["code"] != token_data.get("code"):77            raise EmailCodeError()78 79        return {"is_valid": True, "email": token_data.get("email")}80 81 82class ForgotPasswordResetApi(Resource):83    @setup_required84    def post(self):85        parser = reqparse.RequestParser()86        parser.add_argument("token", type=str, required=True, nullable=False, location="json")87        parser.add_argument("new_password", type=valid_password, required=True, nullable=False, location="json")88        parser.add_argument("password_confirm", type=valid_password, required=True, nullable=False, location="json")89        args = parser.parse_args()90 91        new_password = args["new_password"]92        password_confirm = args["password_confirm"]93 94        if str(new_password).strip() != str(password_confirm).strip():95            raise PasswordMismatchError()96 97        token = args["token"]98        reset_data = AccountService.get_reset_password_data(token)99 100        if reset_data is None:101            raise InvalidTokenError()102 103        AccountService.revoke_reset_password_token(token)104 105        salt = secrets.token_bytes(16)106        base64_salt = base64.b64encode(salt).decode()107 108        password_hashed = hash_password(new_password, salt)109        base64_password_hashed = base64.b64encode(password_hashed).decode()110 111        account = Account.query.filter_by(email=reset_data.get("email")).first()112        if account:113            account.password = base64_password_hashed114            account.password_salt = base64_salt115            db.session.commit()116            tenant = TenantService.get_join_tenants(account)117            if not tenant and not FeatureService.get_system_features().is_allow_create_workspace:118                tenant = TenantService.create_tenant(f"{account.name}'s Workspace")119                TenantService.create_tenant_member(tenant, account, role="owner")120                account.current_tenant = tenant121                tenant_was_created.send(tenant)122        else:123            try:124                account = AccountService.create_account_and_tenant(125                    email=reset_data.get("email"),126                    name=reset_data.get("email"),127                    password=password_confirm,128                    interface_language=languages[0],129                )130            except WorkSpaceNotAllowedCreateError:131                pass132 133        return {"result": "success"}134 135 136api.add_resource(ForgotPasswordSendEmailApi, "/forgot-password")137api.add_resource(ForgotPasswordCheckApi, "/forgot-password/validity")138api.add_resource(ForgotPasswordResetApi, "/forgot-password/resets")139