feat: add init login type

langgenius · Oct 17, 2024 · be4b62d · be4b62d
1 parent abe6a1b
commit be4b62d
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 3 deletions.
diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py
@@ -15,7 +15,12 @@
     InvalidEmailError,
     InvalidTokenError,
 )
-from controllers.console.error import EmailSendIpLimitError, NotAllowedCreateWorkspace, NotAllowedRegister
+from controllers.console.error import (
+    AccountBannedOrClosedError,
+    EmailSendIpLimitError,
+    NotAllowedCreateWorkspace,
+    NotAllowedRegister,
+)
 from controllers.console.setup import setup_required
 from events.tenant_event import tenant_was_created
 from libs.helper import email, extract_remote_ip
@@ -62,7 +67,7 @@ def post(self):
             else:
                 account = AccountService.authenticate(args["email"], args["password"])
         except services.errors.account.AccountLoginError:
-            raise NotAllowedRegister()
+            raise AccountBannedOrClosedError()
         except services.errors.account.AccountPasswordError:
             AccountService.add_login_error_rate_limit(args["email"])
             raise EmailOrPasswordMismatchError()

diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py
@@ -46,9 +46,15 @@ class NotAllowedCreateWorkspace(BaseHTTPException):
     code = 400
 
 
+class AccountBannedOrClosedError(BaseHTTPException):
+    error_code = "account_banned_or_closed"
+    description = "Account is banned or closed."
+    code = 400
+
+
 class NotAllowedRegister(BaseHTTPException):
     error_code = "unauthorized"
-    description = "Account not found."
+    description = "account_not_found"
     code = 400
 
 

diff --git a/api/services/account_service.py b/api/services/account_service.py
@@ -32,6 +32,7 @@
     TenantStatus,
 )
 from models.model import DifySetup
+from services.enterprise.enterprise_service import EnterpriseService
 from services.errors.account import (
     AccountAlreadyInTenantError,
     AccountLoginError,
@@ -953,3 +954,13 @@ def _get_invitation_by_token(
 def _generate_refresh_token(length: int = 64):
     token = secrets.token_hex(length)
     return token
+
+
+class InitLoginType:
+    @staticmethod
+    def init_login_type():
+        enterprise_info = EnterpriseService.get_info()
+        dify_config.ENABLE_EMAIL_CODE_LOGIN = enterprise_info["enable_email_code_login"]
+        dify_config.ENABLE_EMAIL_PASSWORD_LOGIN = enterprise_info["enable_email_password_login"]
+        dify_config.ALLOW_REGISTER = enterprise_info["is_allow_register"]
+        dify_config.ALLOW_CREATE_WORKSPACE = enterprise_info["is_allow_create_workspace"]
diff --git a/api/services/feature_service.py b/api/services/feature_service.py
@@ -1,6 +1,7 @@
 from pydantic import BaseModel, ConfigDict
 
 from configs import dify_config
+from services.account_service import InitLoginType
 from services.billing_service import BillingService
 from services.enterprise.enterprise_service import EnterpriseService
 
@@ -126,6 +127,7 @@ def _fulfill_params_from_billing_api(cls, features: FeatureModel, tenant_id: str
 
     @classmethod
     def _fulfill_params_from_enterprise(cls, features):
+        InitLoginType.init_login_type()
         enterprise_info = EnterpriseService.get_info()
 
         features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"]