Underground-Digital/Workflow-Engine
0
1from typing import cast2 3import flask_login4from flask import request5from flask_restful import Resource, reqparse6 7import services8from constants.languages import languages9from controllers.console import api10from controllers.console.auth.error import (11 EmailCodeError,12 EmailOrPasswordMismatchError,13 EmailPasswordLoginLimitError,14 InvalidEmailError,15 InvalidTokenError,16)17from controllers.console.error import (18 AccountBannedError,19 EmailSendIpLimitError,20 NotAllowedCreateWorkspace,21 NotAllowedRegister,22)23from controllers.console.wraps import setup_required24from events.tenant_event import tenant_was_created25from libs.helper import email, extract_remote_ip26from libs.password import valid_password27from models.account import Account28from services.account_service import AccountService, RegisterService, TenantService29from services.errors.workspace import WorkSpaceNotAllowedCreateError30from services.feature_service import FeatureService31 32 33class LoginApi(Resource):34 """Resource for user login."""35 36 @setup_required37 def post(self):38 """Authenticate user and login."""39 parser = reqparse.RequestParser()40 parser.add_argument("email", type=email, required=True, location="json")41 parser.add_argument("password", type=valid_password, required=True, location="json")42 parser.add_argument("remember_me", type=bool, required=False, default=False, location="json")43 parser.add_argument("invite_token", type=str, required=False, default=None, location="json")44 parser.add_argument("language", type=str, required=False, default="en-US", location="json")45 args = parser.parse_args()46 47 is_login_error_rate_limit = AccountService.is_login_error_rate_limit(args["email"])48 if is_login_error_rate_limit:49 raise EmailPasswordLoginLimitError()50 51 invitation = args["invite_token"]52 if invitation:53 invitation = RegisterService.get_invitation_if_token_valid(None, args["email"], invitation)54 55 if args["language"] is not None and args["language"] == "zh-Hans":56 language = "zh-Hans"57 else:58 language = "en-US"59 60 try:61 if invitation:62 data = invitation.get("data", {})63 invitee_email = data.get("email") if data else None64 if invitee_email != args["email"]:65 raise InvalidEmailError()66 account = AccountService.authenticate(args["email"], args["password"], args["invite_token"])67 else:68 account = AccountService.authenticate(args["email"], args["password"])69 except services.errors.account.AccountLoginError:70 raise AccountBannedError()71 except services.errors.account.AccountPasswordError:72 AccountService.add_login_error_rate_limit(args["email"])73 raise EmailOrPasswordMismatchError()74 except services.errors.account.AccountNotFoundError:75 if FeatureService.get_system_features().is_allow_register:76 token = AccountService.send_reset_password_email(email=args["email"], language=language)77 return {"result": "fail", "data": token, "code": "account_not_found"}78 else:79 raise NotAllowedRegister()80 # SELF_HOSTED only have one workspace81 tenants = TenantService.get_join_tenants(account)82 if len(tenants) == 0:83 return {84 "result": "fail",85 "data": "workspace not found, please contact system admin to invite you to join in a workspace",86 }87 88 token_pair = AccountService.login(account=account, ip_address=extract_remote_ip(request))89 AccountService.reset_login_error_rate_limit(args["email"])90 return {"result": "success", "data": token_pair.model_dump()}91 92 93class LogoutApi(Resource):94 @setup_required95 def get(self):96 account = cast(Account, flask_login.current_user)97 if isinstance(account, flask_login.AnonymousUserMixin):98 return {"result": "success"}99 AccountService.logout(account=account)100 flask_login.logout_user()101 return {"result": "success"}102 103 104class ResetPasswordSendEmailApi(Resource):105 @setup_required106 def post(self):107 parser = reqparse.RequestParser()108 parser.add_argument("email", type=email, required=True, location="json")109 parser.add_argument("language", type=str, required=False, location="json")110 args = parser.parse_args()111 112 if args["language"] is not None and args["language"] == "zh-Hans":113 language = "zh-Hans"114 else:115 language = "en-US"116 117 account = AccountService.get_user_through_email(args["email"])118 if account is None:119 if FeatureService.get_system_features().is_allow_register:120 token = AccountService.send_reset_password_email(email=args["email"], language=language)121 else:122 raise NotAllowedRegister()123 else:124 token = AccountService.send_reset_password_email(account=account, language=language)125 126 return {"result": "success", "data": token}127 128 129class EmailCodeLoginSendEmailApi(Resource):130 @setup_required131 def post(self):132 parser = reqparse.RequestParser()133 parser.add_argument("email", type=email, required=True, location="json")134 parser.add_argument("language", type=str, required=False, location="json")135 args = parser.parse_args()136 137 ip_address = extract_remote_ip(request)138 if AccountService.is_email_send_ip_limit(ip_address):139 raise EmailSendIpLimitError()140 141 if args["language"] is not None and args["language"] == "zh-Hans":142 language = "zh-Hans"143 else:144 language = "en-US"145 146 account = AccountService.get_user_through_email(args["email"])147 if account is None:148 if FeatureService.get_system_features().is_allow_register:149 token = AccountService.send_email_code_login_email(email=args["email"], language=language)150 else:151 raise NotAllowedRegister()152 else:153 token = AccountService.send_email_code_login_email(account=account, language=language)154 155 return {"result": "success", "data": token}156 157 158class EmailCodeLoginApi(Resource):159 @setup_required160 def post(self):161 parser = reqparse.RequestParser()162 parser.add_argument("email", type=str, required=True, location="json")163 parser.add_argument("code", type=str, required=True, location="json")164 parser.add_argument("token", type=str, required=True, location="json")165 args = parser.parse_args()166 167 user_email = args["email"]168 169 token_data = AccountService.get_email_code_login_data(args["token"])170 if token_data is None:171 raise InvalidTokenError()172 173 if token_data["email"] != args["email"]:174 raise InvalidEmailError()175 176 if token_data["code"] != args["code"]:177 raise EmailCodeError()178 179 AccountService.revoke_email_code_login_token(args["token"])180 account = AccountService.get_user_through_email(user_email)181 if account:182 tenant = TenantService.get_join_tenants(account)183 if not tenant:184 if not FeatureService.get_system_features().is_allow_create_workspace:185 raise NotAllowedCreateWorkspace()186 else:187 tenant = TenantService.create_tenant(f"{account.name}'s Workspace")188 TenantService.create_tenant_member(tenant, account, role="owner")189 account.current_tenant = tenant190 tenant_was_created.send(tenant)191 192 if account is None:193 try:194 account = AccountService.create_account_and_tenant(195 email=user_email, name=user_email, interface_language=languages[0]196 )197 except WorkSpaceNotAllowedCreateError:198 return NotAllowedCreateWorkspace()199 token_pair = AccountService.login(account, ip_address=extract_remote_ip(request))200 AccountService.reset_login_error_rate_limit(args["email"])201 return {"result": "success", "data": token_pair.model_dump()}202 203 204class RefreshTokenApi(Resource):205 def post(self):206 parser = reqparse.RequestParser()207 parser.add_argument("refresh_token", type=str, required=True, location="json")208 args = parser.parse_args()209 210 try:211 new_token_pair = AccountService.refresh_token(args["refresh_token"])212 return {"result": "success", "data": new_token_pair.model_dump()}213 except Exception as e:214 return {"result": "fail", "data": str(e)}, 401215 216 217api.add_resource(LoginApi, "/login")218api.add_resource(LogoutApi, "/logout")219api.add_resource(EmailCodeLoginSendEmailApi, "/email-code-login")220api.add_resource(EmailCodeLoginApi, "/email-code-login/validity")221api.add_resource(ResetPasswordSendEmailApi, "/reset-password")222api.add_resource(RefreshTokenApi, "/refresh-token")223