Dear User,
- #The Dify team has generated a new password for you, details as follows:
- #{new_password}
- #Please change your password to log in as soon as possible.
- #Regards,
- #The Dify Team
- # """ - # } - # response = mailchimp.messages.send({ - # 'message': message, - # # required for transactional email - # ' settings': { - # 'sandbox_mode': dify_config.MAILCHIMP_SANDBOX_MODE, - # }, - # }) +class EmailCodeLoginSendEmailApi(Resource): + @setup_required + def post(self): + parser = reqparse.RequestParser() + parser.add_argument("email", type=email, required=True, location="json") + parser.add_argument("language", type=str, required=False, location="json") + args = parser.parse_args() - # Check if MSG was sent - # if response.status_code != 200: - # # handle error - # pass + ip_address = extract_remote_ip(request) + if AccountService.is_email_send_ip_limit(ip_address): + raise EmailSendIpLimitError() - return {"result": "success"} + if args["language"] is not None and args["language"] == "zh-Hans": + language = "zh-Hans" + else: + language = "en-US" + + account = AccountService.get_user_through_email(args["email"]) + if account is None: + if FeatureService.get_system_features().is_allow_register: + token = AccountService.send_email_code_login_email(email=args["email"], language=language) + else: + raise NotAllowedRegister() + else: + token = AccountService.send_email_code_login_email(account=account, language=language) + + return {"result": "success", "data": token} + + +class EmailCodeLoginApi(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, location="json") + args = parser.parse_args() + + user_email = args["email"] + + token_data = AccountService.get_email_code_login_data(args["token"]) + if token_data is None: + raise InvalidTokenError() + + if token_data["email"] != args["email"]: + raise InvalidEmailError() + + if token_data["code"] != args["code"]: + raise EmailCodeError() + + AccountService.revoke_email_code_login_token(args["token"]) + account = AccountService.get_user_through_email(user_email) + if account: + tenant = TenantService.get_join_tenants(account) + if not tenant: + if not FeatureService.get_system_features().is_allow_create_workspace: + raise NotAllowedCreateWorkspace() + else: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role="owner") + account.current_tenant = tenant + tenant_was_created.send(tenant) + + if account is None: + try: + account = AccountService.create_account_and_tenant( + email=user_email, name=user_email, interface_language=languages[0] + ) + except WorkSpaceNotAllowedCreateError: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin" + "?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) + token_pair = AccountService.login(account, ip_address=extract_remote_ip(request)) + AccountService.reset_login_error_rate_limit(args["email"]) + return {"result": "success", "data": token_pair.model_dump()} class RefreshTokenApi(Resource): @@ -120,4 +220,7 @@ class RefreshTokenApi(Resource): 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(ResetPasswordSendEmailApi, "/reset-password") api.add_resource(RefreshTokenApi, "/refresh-token") diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index c5909b8c10..ba0e07cd16 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -5,14 +5,19 @@ from typing import Optional import requests from flask import current_app, redirect, request from flask_restful import Resource +from werkzeug.exceptions import Unauthorized from configs import dify_config from constants.languages import languages +from events.tenant_event import tenant_was_created from extensions.ext_database import db from libs.helper import extract_remote_ip from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo from models.account import Account, AccountStatus from services.account_service import AccountService, RegisterService, TenantService +from services.errors.account import AccountNotFoundError +from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError +from services.feature_service import FeatureService from .. import api @@ -42,6 +47,7 @@ def get_oauth_providers(): class OAuthLogin(Resource): def get(self, provider: str): + invite_token = request.args.get("invite_token") or None OAUTH_PROVIDERS = get_oauth_providers() with current_app.app_context(): oauth_provider = OAUTH_PROVIDERS.get(provider) @@ -49,7 +55,7 @@ class OAuthLogin(Resource): if not oauth_provider: return {"error": "Invalid provider"}, 400 - auth_url = oauth_provider.get_authorization_url() + auth_url = oauth_provider.get_authorization_url(invite_token=invite_token) return redirect(auth_url) @@ -62,6 +68,11 @@ class OAuthCallback(Resource): return {"error": "Invalid provider"}, 400 code = request.args.get("code") + state = request.args.get("state") + invite_token = None + if state: + invite_token = state + try: token = oauth_provider.get_access_token(code) user_info = oauth_provider.get_user_info(token) @@ -69,7 +80,27 @@ class OAuthCallback(Resource): logging.exception(f"An error occurred during the OAuth process with {provider}: {e.response.text}") return {"error": "OAuth process failed"}, 400 - account = _generate_account(provider, user_info) + if invite_token and RegisterService.is_valid_invite_token(invite_token): + invitation = RegisterService._get_invitation_by_token(token=invite_token) + if invitation: + invitation_email = invitation.get("email", None) + if invitation_email != user_info.email: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Invalid invitation token.") + + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin/invite-settings?invite_token={invite_token}") + + try: + account = _generate_account(provider, user_info) + except AccountNotFoundError: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Account not found.") + except WorkSpaceNotFoundError: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found.") + except WorkSpaceNotAllowedCreateError: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin" + "?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) + # Check account status if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}: return {"error": "Account is banned or closed."}, 403 @@ -79,7 +110,15 @@ class OAuthCallback(Resource): account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None) db.session.commit() - TenantService.create_owner_tenant_if_not_exist(account) + try: + TenantService.create_owner_tenant_if_not_exist(account) + except Unauthorized: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found.") + except WorkSpaceNotAllowedCreateError: + return redirect( + f"{dify_config.CONSOLE_WEB_URL}/signin" + "?message=Workspace not found, please contact system admin to invite you to join in a workspace." + ) token_pair = AccountService.login( account=account, @@ -104,8 +143,20 @@ def _generate_account(provider: str, user_info: OAuthUserInfo): # Get account by openid or email. account = _get_account_by_openid_or_email(provider, user_info) + if account: + tenant = TenantService.get_join_tenants(account) + if not tenant: + if not FeatureService.get_system_features().is_allow_create_workspace: + raise WorkSpaceNotAllowedCreateError() + else: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role="owner") + account.current_tenant = tenant + tenant_was_created.send(tenant) + if not account: - # Create account + if not FeatureService.get_system_features().is_allow_register: + raise AccountNotFoundError() account_name = user_info.name or "Dify" account = RegisterService.register( email=user_info.email, name=account_name, password=None, open_id=user_info.id, provider=provider diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py index 870e547728..a6d4c8e8ec 100644 --- a/api/controllers/console/error.py +++ b/api/controllers/console/error.py @@ -38,3 +38,27 @@ class AlreadyActivateError(BaseHTTPException): error_code = "already_activate" description = "Auth Token is invalid or account already activated, please check again." code = 403 + + +class NotAllowedCreateWorkspace(BaseHTTPException): + error_code = "unauthorized" + description = "Workspace not found, please contact system admin to invite you to join in a workspace." + code = 400 + + +class AccountBannedError(BaseHTTPException): + error_code = "account_banned" + description = "Account is banned." + code = 400 + + +class NotAllowedRegister(BaseHTTPException): + error_code = "unauthorized" + description = "Account not found." + code = 400 + + +class EmailSendIpLimitError(BaseHTTPException): + error_code = "email_send_ip_limit" + description = "Too many emails have been sent from this IP address recently. Please try again later." + code = 429 diff --git a/api/libs/helper.py b/api/libs/helper.py index d8a8e7f411..fc3be8473c 100644 --- a/api/libs/helper.py +++ b/api/libs/helper.py @@ -189,23 +189,39 @@ def compact_generate_response(response: Union[dict, RateLimitGenerator]) -> Resp class TokenManager: @classmethod - def generate_token(cls, account: Account, token_type: str, additional_data: Optional[dict] = None) -> str: - old_token = cls._get_current_token_for_account(account.id, token_type) - if old_token: - if isinstance(old_token, bytes): - old_token = old_token.decode("utf-8") - cls.revoke_token(old_token, token_type) + def generate_token( + cls, + token_type: str, + account: Optional[Account] = None, + email: Optional[str] = None, + additional_data: Optional[dict] = None, + ) -> str: + if account is None and email is None: + raise ValueError("Account or email must be provided") + + account_id = account.id if account else None + account_email = account.email if account else email + + if account_id: + old_token = cls._get_current_token_for_account(account_id, token_type) + if old_token: + if isinstance(old_token, bytes): + old_token = old_token.decode("utf-8") + cls.revoke_token(old_token, token_type) token = str(uuid.uuid4()) - token_data = {"account_id": account.id, "email": account.email, "token_type": token_type} + token_data = {"account_id": account_id, "email": account_email, "token_type": token_type} if additional_data: token_data.update(additional_data) expiry_hours = current_app.config[f"{token_type.upper()}_TOKEN_EXPIRY_HOURS"] token_key = cls._get_token_key(token, token_type) - redis_client.setex(token_key, expiry_hours * 60 * 60, json.dumps(token_data)) + expiry_time = int(expiry_hours * 60 * 60) + redis_client.setex(token_key, expiry_time, json.dumps(token_data)) + + if account_id: + cls._set_current_token_for_account(account.id, token, token_type, expiry_hours) - cls._set_current_token_for_account(account.id, token, token_type, expiry_hours) return token @classmethod @@ -234,9 +250,12 @@ class TokenManager: return current_token @classmethod - def _set_current_token_for_account(cls, account_id: str, token: str, token_type: str, expiry_hours: int): + def _set_current_token_for_account( + cls, account_id: str, token: str, token_type: str, expiry_hours: Union[int, float] + ): key = cls._get_account_token_key(account_id, token_type) - redis_client.setex(key, expiry_hours * 60 * 60, token) + expiry_time = int(expiry_hours * 60 * 60) + redis_client.setex(key, expiry_time, token) @classmethod def _get_account_token_key(cls, account_id: str, token_type: str) -> str: diff --git a/api/libs/oauth.py b/api/libs/oauth.py index d8ce1a1e66..6b6919de24 100644 --- a/api/libs/oauth.py +++ b/api/libs/oauth.py @@ -1,5 +1,6 @@ import urllib.parse from dataclasses import dataclass +from typing import Optional import requests @@ -40,12 +41,14 @@ class GitHubOAuth(OAuth): _USER_INFO_URL = "https://api.github.com/user" _EMAIL_INFO_URL = "https://api.github.com/user/emails" - def get_authorization_url(self): + def get_authorization_url(self, invite_token: Optional[str] = None): params = { "client_id": self.client_id, "redirect_uri": self.redirect_uri, "scope": "user:email", # Request only basic user information } + if invite_token: + params["state"] = invite_token return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}" def get_access_token(self, code: str): @@ -90,13 +93,15 @@ class GoogleOAuth(OAuth): _TOKEN_URL = "https://oauth2.googleapis.com/token" _USER_INFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo" - def get_authorization_url(self): + def get_authorization_url(self, invite_token: Optional[str] = None): params = { "client_id": self.client_id, "response_type": "code", "redirect_uri": self.redirect_uri, "scope": "openid email", } + if invite_token: + params["state"] = invite_token return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}" def get_access_token(self, code: str): diff --git a/api/libs/password.py b/api/libs/password.py index cfcc0db22d..cdf55c57e5 100644 --- a/api/libs/password.py +++ b/api/libs/password.py @@ -13,7 +13,7 @@ def valid_password(password): if re.match(pattern, password) is not None: return password - raise ValueError("Not a valid password.") + raise ValueError("Password must contain letters and numbers, and the length must be greater than 8.") def hash_password(password_str, salt_byte): diff --git a/api/services/account_service.py b/api/services/account_service.py index eda6011aef..529b716773 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -1,6 +1,7 @@ import base64 import json import logging +import random import secrets import uuid from datetime import datetime, timedelta, timezone @@ -34,7 +35,9 @@ from models.model import DifySetup from services.errors.account import ( AccountAlreadyInTenantError, AccountLoginError, + AccountNotFoundError, AccountNotLinkTenantError, + AccountPasswordError, AccountRegisterError, CannotOperateSelfError, CurrentPasswordIncorrectError, @@ -42,10 +45,12 @@ from services.errors.account import ( LinkAccountIntegrateError, MemberNotInTenantError, NoPermissionError, - RateLimitExceededError, RoleAlreadyAssignedError, TenantNotFoundError, ) +from services.errors.workspace import WorkSpaceNotAllowedCreateError +from services.feature_service import FeatureService +from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task @@ -61,7 +66,11 @@ REFRESH_TOKEN_EXPIRY = timedelta(days=30) class AccountService: - reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=5, time_window=60 * 60) + reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=1, time_window=60 * 1) + email_code_login_rate_limiter = RateLimiter( + prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1 + ) + LOGIN_MAX_ERROR_LIMITS = 5 @staticmethod def _get_refresh_token_key(refresh_token: str) -> str: @@ -127,23 +136,34 @@ class AccountService: return token @staticmethod - def authenticate(email: str, password: str) -> Account: + def authenticate(email: str, password: str, invite_token: Optional[str] = None) -> Account: """authenticate account with email and password""" account = Account.query.filter_by(email=email).first() if not account: - raise AccountLoginError("Invalid email or password.") + raise AccountNotFoundError() if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}: raise AccountLoginError("Account is banned or closed.") + if password and invite_token and account.password is None: + # if invite_token is valid, set password and password_salt + salt = secrets.token_bytes(16) + base64_salt = base64.b64encode(salt).decode() + password_hashed = hash_password(password, salt) + base64_password_hashed = base64.b64encode(password_hashed).decode() + account.password = base64_password_hashed + account.password_salt = base64_salt + + if account.password is None or not compare_password(password, account.password, account.password_salt): + raise AccountPasswordError("Invalid email or password.") + if account.status == AccountStatus.PENDING.value: account.status = AccountStatus.ACTIVE.value account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None) - db.session.commit() - if account.password is None or not compare_password(password, account.password, account.password_salt): - raise AccountLoginError("Invalid email or password.") + db.session.commit() + return account @staticmethod @@ -169,9 +189,18 @@ class AccountService: @staticmethod def create_account( - email: str, name: str, interface_language: str, password: Optional[str] = None, interface_theme: str = "light" + email: str, + name: str, + interface_language: str, + password: Optional[str] = None, + interface_theme: str = "light", + is_setup: Optional[bool] = False, ) -> Account: """create account""" + if not FeatureService.get_system_features().is_allow_register and not is_setup: + from controllers.console.error import NotAllowedRegister + + raise NotAllowedRegister() account = Account() account.email = email account.name = name @@ -198,6 +227,19 @@ class AccountService: db.session.commit() return account + @staticmethod + def create_account_and_tenant( + email: str, name: str, interface_language: str, password: Optional[str] = None + ) -> Account: + """create account""" + account = AccountService.create_account( + email=email, name=name, interface_language=interface_language, password=password + ) + + TenantService.create_owner_tenant_if_not_exist(account=account) + + return account + @staticmethod def link_account_integrate(provider: str, open_id: str, account: Account) -> None: """Link account integrate""" @@ -256,6 +298,10 @@ class AccountService: if ip_address: AccountService.update_login_info(account=account, ip_address=ip_address) + if account.status == AccountStatus.PENDING.value: + account.status = AccountStatus.ACTIVE.value + db.session.commit() + access_token = AccountService.get_account_jwt_token(account=account) refresh_token = _generate_refresh_token() @@ -294,13 +340,29 @@ class AccountService: return AccountService.load_user(account_id) @classmethod - def send_reset_password_email(cls, account): - 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.") + def send_reset_password_email( + cls, + account: Optional[Account] = None, + email: Optional[str] = None, + language: Optional[str] = "en-US", + ): + account_email = account.email if account else email - token = TokenManager.generate_token(account, "reset_password") - send_reset_password_mail_task.delay(language=account.interface_language, to=account.email, token=token) - cls.reset_password_rate_limiter.increment_rate_limit(account.email) + if cls.reset_password_rate_limiter.is_rate_limited(account_email): + from controllers.console.auth.error import PasswordResetRateLimitExceededError + + raise PasswordResetRateLimitExceededError() + + 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=language, + to=account_email, + code=code, + ) + cls.reset_password_rate_limiter.increment_rate_limit(account_email) return token @classmethod @@ -311,11 +373,125 @@ class AccountService: def get_reset_password_data(cls, token: str) -> Optional[dict[str, Any]]: return TokenManager.get_token_data(token, "reset_password") + @classmethod + def send_email_code_login_email( + cls, account: Optional[Account] = None, email: Optional[str] = None, language: Optional[str] = "en-US" + ): + if cls.email_code_login_rate_limiter.is_rate_limited(email): + from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError + + raise EmailCodeLoginRateLimitExceededError() + + code = "".join([str(random.randint(0, 9)) for _ in range(6)]) + token = TokenManager.generate_token( + account=account, email=email, token_type="email_code_login", additional_data={"code": code} + ) + send_email_code_login_mail_task.delay( + language=language, + to=account.email if account else email, + code=code, + ) + cls.email_code_login_rate_limiter.increment_rate_limit(email) + return token + + @classmethod + def get_email_code_login_data(cls, token: str) -> Optional[dict[str, Any]]: + return TokenManager.get_token_data(token, "email_code_login") + + @classmethod + def revoke_email_code_login_token(cls, token: str): + TokenManager.revoke_token(token, "email_code_login") + + @classmethod + def get_user_through_email(cls, email: str): + account = db.session.query(Account).filter(Account.email == email).first() + if not account: + return None + + if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}: + raise Unauthorized("Account is banned or closed.") + + return account + + @staticmethod + def add_login_error_rate_limit(email: str) -> None: + key = f"login_error_rate_limit:{email}" + count = redis_client.get(key) + if count is None: + count = 0 + count = int(count) + 1 + redis_client.setex(key, 60 * 60 * 24, count) + + @staticmethod + def is_login_error_rate_limit(email: str) -> bool: + key = f"login_error_rate_limit:{email}" + count = redis_client.get(key) + if count is None: + return False + + count = int(count) + if count > AccountService.LOGIN_MAX_ERROR_LIMITS: + return True + return False + + @staticmethod + def reset_login_error_rate_limit(email: str): + key = f"login_error_rate_limit:{email}" + redis_client.delete(key) + + @staticmethod + def is_email_send_ip_limit(ip_address: str): + minute_key = f"email_send_ip_limit_minute:{ip_address}" + freeze_key = f"email_send_ip_limit_freeze:{ip_address}" + hour_limit_key = f"email_send_ip_limit_hour:{ip_address}" + + # check ip is frozen + if redis_client.get(freeze_key): + return True + + # check current minute count + current_minute_count = redis_client.get(minute_key) + if current_minute_count is None: + current_minute_count = 0 + current_minute_count = int(current_minute_count) + + # check current hour count + if current_minute_count > dify_config.EMAIL_SEND_IP_LIMIT_PER_MINUTE: + hour_limit_count = redis_client.get(hour_limit_key) + if hour_limit_count is None: + hour_limit_count = 0 + hour_limit_count = int(hour_limit_count) + + if hour_limit_count >= 1: + redis_client.setex(freeze_key, 60 * 60, 1) + return True + else: + redis_client.setex(hour_limit_key, 60 * 10, hour_limit_count + 1) # first time limit 10 minutes + + # add hour limit count + redis_client.incr(hour_limit_key) + redis_client.expire(hour_limit_key, 60 * 60) + + return True + + redis_client.setex(minute_key, 60, current_minute_count + 1) + redis_client.expire(minute_key, 60) + + return False + + +def _get_login_cache_key(*, account_id: str, token: str): + return f"account_login:{account_id}:{token}" + class TenantService: @staticmethod - def create_tenant(name: str) -> Tenant: + def create_tenant(name: str, is_setup: Optional[bool] = False) -> Tenant: """Create tenant""" + if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup: + from controllers.console.error import NotAllowedCreateWorkspace + + raise NotAllowedCreateWorkspace() tenant = Tenant(name=name) db.session.add(tenant) @@ -326,8 +502,12 @@ class TenantService: return tenant @staticmethod - def create_owner_tenant_if_not_exist(account: Account, name: Optional[str] = None): + def create_owner_tenant_if_not_exist( + account: Account, name: Optional[str] = None, is_setup: Optional[bool] = False + ): """Create owner tenant if not exist""" + if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup: + raise WorkSpaceNotAllowedCreateError() available_ta = ( TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first() ) @@ -336,9 +516,9 @@ class TenantService: return if name: - tenant = TenantService.create_tenant(name) + tenant = TenantService.create_tenant(name=name, is_setup=is_setup) else: - tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup) TenantService.create_tenant_member(tenant, account, role="owner") account.current_tenant = tenant db.session.commit() @@ -352,8 +532,13 @@ class TenantService: logging.error(f"Tenant {tenant.id} has already an owner.") raise Exception("Tenant already has an owner.") - ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role) - db.session.add(ta) + ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first() + if ta: + ta.role = role + else: + ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role) + db.session.add(ta) + db.session.commit() return ta @@ -570,12 +755,13 @@ class RegisterService: name=name, interface_language=languages[0], password=password, + is_setup=True, ) account.last_login_ip = ip_address account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None) - TenantService.create_owner_tenant_if_not_exist(account) + TenantService.create_owner_tenant_if_not_exist(account=account, is_setup=True) dify_setup = DifySetup(version=dify_config.CURRENT_VERSION) db.session.add(dify_setup) @@ -600,27 +786,33 @@ class RegisterService: provider: Optional[str] = None, language: Optional[str] = None, status: Optional[AccountStatus] = None, + is_setup: Optional[bool] = False, ) -> Account: db.session.begin_nested() """Register account""" try: account = AccountService.create_account( - email=email, name=name, interface_language=language or languages[0], password=password + email=email, + name=name, + interface_language=language or languages[0], + password=password, + is_setup=is_setup, ) account.status = AccountStatus.ACTIVE.value if not status else status.value account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None) if open_id is not None or provider is not None: AccountService.link_account_integrate(provider, open_id, account) - if dify_config.EDITION != "SELF_HOSTED": - tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + if FeatureService.get_system_features().is_allow_create_workspace: + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") TenantService.create_tenant_member(tenant, account, role="owner") account.current_tenant = tenant - tenant_was_created.send(tenant) db.session.commit() + except WorkSpaceNotAllowedCreateError: + db.session.rollback() except Exception as e: db.session.rollback() logging.error(f"Register failed: {e}") @@ -639,7 +831,9 @@ class RegisterService: TenantService.check_member_permission(tenant, inviter, None, "add") name = email.split("@")[0] - account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING) + account = cls.register( + email=email, name=name, language=language, status=AccountStatus.PENDING, is_setup=True + ) # Create new tenant member for invited tenant TenantService.create_tenant_member(tenant, account, role) TenantService.switch_tenant(account, tenant.id) @@ -679,6 +873,11 @@ class RegisterService: redis_client.setex(cls._get_invitation_token_key(token), expiry_hours * 60 * 60, json.dumps(invitation_data)) return token + @classmethod + def is_valid_invite_token(cls, token: str) -> bool: + data = redis_client.get(cls._get_invitation_token_key(token)) + return data is not None + @classmethod def revoke_token(cls, workspace_id: str, email: str, token: str): if workspace_id and email: @@ -727,7 +926,9 @@ class RegisterService: } @classmethod - def _get_invitation_by_token(cls, token: str, workspace_id: str, email: str) -> Optional[dict[str, str]]: + def _get_invitation_by_token( + cls, token: str, workspace_id: Optional[str] = None, email: Optional[str] = None + ) -> Optional[dict[str, str]]: if workspace_id is not None and email is not None: email_hash = sha256(email.encode()).hexdigest() cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}" diff --git a/api/services/errors/account.py b/api/services/errors/account.py index 82dd9f944a..5aca12ffeb 100644 --- a/api/services/errors/account.py +++ b/api/services/errors/account.py @@ -13,6 +13,10 @@ class AccountLoginError(BaseServiceError): pass +class AccountPasswordError(BaseServiceError): + pass + + class AccountNotLinkTenantError(BaseServiceError): pass diff --git a/api/services/errors/workspace.py b/api/services/errors/workspace.py new file mode 100644 index 0000000000..714064ffdf --- /dev/null +++ b/api/services/errors/workspace.py @@ -0,0 +1,9 @@ +from services.errors.base import BaseServiceError + + +class WorkSpaceNotAllowedCreateError(BaseServiceError): + pass + + +class WorkSpaceNotFoundError(BaseServiceError): + pass diff --git a/api/services/feature_service.py b/api/services/feature_service.py index 4d5812c6c6..c321393bc5 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -42,6 +42,11 @@ class SystemFeatureModel(BaseModel): sso_enforced_for_web: bool = False sso_enforced_for_web_protocol: str = "" enable_web_sso_switch_component: bool = False + enable_email_code_login: bool = False + enable_email_password_login: bool = True + enable_social_oauth_login: bool = False + is_allow_register: bool = False + is_allow_create_workspace: bool = False class FeatureService: @@ -60,12 +65,23 @@ class FeatureService: def get_system_features(cls) -> SystemFeatureModel: system_features = SystemFeatureModel() + cls._fulfill_system_params_from_env(system_features) + if dify_config.ENTERPRISE_ENABLED: system_features.enable_web_sso_switch_component = True + cls._fulfill_params_from_enterprise(system_features) return system_features + @classmethod + def _fulfill_system_params_from_env(cls, system_features: SystemFeatureModel): + system_features.enable_email_code_login = dify_config.ENABLE_EMAIL_CODE_LOGIN + system_features.enable_email_password_login = dify_config.ENABLE_EMAIL_PASSWORD_LOGIN + system_features.enable_social_oauth_login = dify_config.ENABLE_SOCIAL_OAUTH_LOGIN + system_features.is_allow_register = dify_config.ALLOW_REGISTER + system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE + @classmethod def _fulfill_params_from_env(cls, features: FeatureModel): features.can_replace_logo = dify_config.CAN_REPLACE_LOGO @@ -113,7 +129,19 @@ class FeatureService: def _fulfill_params_from_enterprise(cls, features): enterprise_info = EnterpriseService.get_info() - features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"] - features.sso_enforced_for_signin_protocol = enterprise_info["sso_enforced_for_signin_protocol"] - features.sso_enforced_for_web = enterprise_info["sso_enforced_for_web"] - features.sso_enforced_for_web_protocol = enterprise_info["sso_enforced_for_web_protocol"] + if "sso_enforced_for_signin" in enterprise_info: + features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"] + if "sso_enforced_for_signin_protocol" in enterprise_info: + features.sso_enforced_for_signin_protocol = enterprise_info["sso_enforced_for_signin_protocol"] + if "sso_enforced_for_web" in enterprise_info: + features.sso_enforced_for_web = enterprise_info["sso_enforced_for_web"] + if "sso_enforced_for_web_protocol" in enterprise_info: + features.sso_enforced_for_web_protocol = enterprise_info["sso_enforced_for_web_protocol"] + if "enable_email_code_login" in enterprise_info: + features.enable_email_code_login = enterprise_info["enable_email_code_login"] + if "enable_email_password_login" in enterprise_info: + features.enable_email_password_login = enterprise_info["enable_email_password_login"] + if "is_allow_register" in enterprise_info: + features.is_allow_register = enterprise_info["is_allow_register"] + if "is_allow_create_workspace" in enterprise_info: + features.is_allow_create_workspace = enterprise_info["is_allow_create_workspace"] diff --git a/api/tasks/mail_email_code_login.py b/api/tasks/mail_email_code_login.py new file mode 100644 index 0000000000..d78fc2b891 --- /dev/null +++ b/api/tasks/mail_email_code_login.py @@ -0,0 +1,41 @@ +import logging +import time + +import click +from celery import shared_task +from flask import render_template + +from extensions.ext_mail import mail + + +@shared_task(queue="mail") +def send_email_code_login_mail_task(language: str, to: str, code: str): + """ + Async Send email code login mail + :param language: Language in which the email should be sent (e.g., 'en', 'zh') + :param to: Recipient email address + :param code: Email code to be included in the email + """ + if not mail.is_inited(): + return + + logging.info(click.style("Start email code login mail to {}".format(to), fg="green")) + start_at = time.perf_counter() + + # send email code login mail using different languages + try: + if language == "zh-Hans": + html_content = render_template("email_code_login_mail_template_zh-CN.html", to=to, code=code) + mail.send(to=to, subject="邮箱验证码", html=html_content) + else: + html_content = render_template("email_code_login_mail_template_en-US.html", to=to, code=code) + mail.send(to=to, subject="Email Code", html=html_content) + + end_at = time.perf_counter() + logging.info( + click.style( + "Send email code login mail to {} succeeded: latency: {}".format(to, end_at - start_at), fg="green" + ) + ) + except Exception: + logging.exception("Send email code login mail to {} failed".format(to)) diff --git a/api/tasks/mail_reset_password_task.py b/api/tasks/mail_reset_password_task.py index cbb78976ca..8596ca07cf 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,13 +24,12 @@ 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) - mail.send(to=to, subject="重置您的 Dify 密码", html=html_content) + 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) - mail.send(to=to, subject="Reset Your Dify Password", html=html_content) + html_content = render_template("reset_password_mail_template_en-US.html", to=to, code=code) + mail.send(to=to, subject="Set Your Dify Password", html=html_content) end_at = time.perf_counter() logging.info( diff --git a/api/templates/email_code_login_mail_template_en-US.html b/api/templates/email_code_login_mail_template_en-US.html new file mode 100644 index 0000000000..066818d10c --- /dev/null +++ b/api/templates/email_code_login_mail_template_en-US.html @@ -0,0 +1,74 @@ + + + + + + +
+ Your login code for Dify
+Copy and paste this code, this code will only be valid for the next 5 minutes.
+If you didn't request a login, don't worry. You can safely ignore this email.
+
+ Dify 的登录验证码
+复制并粘贴此验证码，注意验证码仅在接下来的 5 分钟内有效。
+如果您没有请求登录，请不要担心。您可以安全地忽略此电子邮件。
+Dear {{ to }},
{{ inviter_name }} is pleased to invite you to join our workspace on Dify, a platform specifically designed for LLM application development. On Dify, you can explore, create, and collaborate to build and operate AI applications.
-You can now log in to Dify using the GitHub or Google account associated with this email.
+Click the button below to log in to Dify and join the workspace.
尊敬的 {{ to }}，
{{ inviter_name }} 现邀请您加入我们在 Dify 的工作区，这是一个专为 LLM 应用开发而设计的平台。在 Dify 上，您可以探索、创造和合作，构建和运营 AI 应用。
-您现在可以使用与此邮件相对应的 GitHub 或 Google 账号登录 Dify。
+点击下方按钮即可登录 Dify 并且加入空间。
- Dear {{ to }},
-We have received a request to reset your password. If you initiated this request, please click the button below to reset your password:
- -If you did not request a password reset, please ignore this email and your account will remain secure.
-
+ Set your Dify password
+Copy and paste this code, this code will only be valid for the next 5 minutes.
+If you didn't request, don't worry. You can safely ignore this email.
- 尊敬的 {{ to }}，
-我们收到了您关于重置密码的请求。如果是您本人操作，请点击以下按钮重置您的密码：
- -如果您没有请求重置密码，请忽略此邮件，您的账户信息将保持安全。
-
+ 设置您的 Dify 账户密码
+复制并粘贴此验证码，注意验证码仅在接下来的 5 分钟内有效。
+如果您没有请求，请不要担心。您可以安全地忽略此电子邮件。