From 6cb0287069150d2f63e62fd411dc218431ebe056 Mon Sep 17 00:00:00 2001
From: ytqh <yutianqiuhao@gmail.com>
Date: Sat, 1 Mar 2025 20:34:05 +0800
Subject: [PATCH] feat: finish send with login in email

---
 .../service_api_with_auth/auth/login.py       | 38 ++++++++-----------
 api/services/account_service.py               | 17 +++++++--
 2 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/api/controllers/service_api_with_auth/auth/login.py b/api/controllers/service_api_with_auth/auth/login.py
index 8edc52bb4d..229287f380 100644
--- a/api/controllers/service_api_with_auth/auth/login.py
+++ b/api/controllers/service_api_with_auth/auth/login.py
@@ -175,7 +175,13 @@ class EmailCodeLoginApi(Resource):
         """
         parser = reqparse.RequestParser()
         # TODO: ytqh add a new field for different tenant (default: Saier)
-        parser.add_argument("tenant_id", type=str, required=False, location="json")
+        parser.add_argument(
+            "tenant_id",
+            type=str,
+            required=False,
+            default="5cd3029e-7f92-428a-a5c8-14a790c70233",
+            location="json",
+        )  # TODO: ytqh move this to the config
         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")
@@ -184,29 +190,17 @@ class EmailCodeLoginApi(Resource):
         user_email = args["email"]
         tenant_id = args["tenant_id"]
 
-        # Skip token validation if in debug mode is True
-        # WARNING: This is for development purposes only and should never be enabled in production
-        if DeploymentConfig().DEBUG:
-            import logging
+        token_data = AccountService.get_email_code_login_data(args["token"])
+        if token_data is None:
+            raise InvalidTokenError()
 
-            logger = logging.getLogger(__name__)
-            logger.warning(
-                f"⚠️ DEBUG MODE: Token validation bypassed for email: {user_email}"
-            )
-
-            tenant_id = "5cd3029e-7f92-428a-a5c8-14a790c70233"  # TODO: ytqh move this to the config
-        else:
-            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["email"] != args["email"]:
+            raise InvalidEmailError()
 
-            if token_data["code"] != args["code"]:
-                raise EmailCodeError()
+        if token_data["code"] != args["code"]:
+            raise EmailCodeError()
 
-            AccountService.revoke_email_code_login_token(args["token"])
+        AccountService.revoke_email_code_login_token(args["token"])
 
         try:
             account = AccountService.get_user_through_email(user_email)
@@ -229,7 +223,7 @@ class EmailCodeLoginApi(Resource):
                 raise AccountInFreezeError()
         else:
             connected_tenant = TenantService.get_join_tenants(account)
-            if connected_tenant is None:
+            if connected_tenant is None or tenant not in connected_tenant:
                 TenantService.create_tenant_member(tenant, account, role="end_user")
 
         token_pair = AccountService.login(
diff --git a/api/services/account_service.py b/api/services/account_service.py
index 177b2c5fae..b1d339ccee 100644
--- a/api/services/account_service.py
+++ b/api/services/account_service.py
@@ -9,6 +9,7 @@ from hashlib import sha256
 from typing import Any, Optional, cast
 
 from configs import dify_config
+from configs.deploy import DeploymentConfig
 from constants.languages import language_timezone_mapping, languages
 from events.tenant_event import tenant_was_created
 from extensions.ext_database import db
@@ -28,7 +29,7 @@ from models.account import (
     TenantStatus,
 )
 from models.model import DifySetup
-from pydantic import BaseModel
+from pydantic import BaseModel  # type: ignore
 from services.billing_service import BillingService
 from services.errors.account import (
     AccountAlreadyInTenantError,
@@ -48,7 +49,7 @@ from services.errors.account import (
 )
 from services.errors.workspace import WorkSpaceNotAllowedCreateError
 from services.feature_service import FeatureService
-from sqlalchemy import func
+from sqlalchemy import func  # type: ignore
 from tasks.delete_account_task import delete_account_task
 from tasks.mail_account_deletion_task import send_account_deletion_verification_code
 from tasks.mail_email_code_login import send_email_code_login_mail_task
@@ -504,14 +505,22 @@ class AccountService:
         email = account.email if account else email
         if email is None:
             raise ValueError("Email must be provided.")
-        if cls.email_code_login_rate_limiter.is_rate_limited(email):
+        if (
+            cls.email_code_login_rate_limiter.is_rate_limited(email)
+            and not DeploymentConfig().DEBUG
+        ):
             from controllers.console.auth.error import (
                 EmailCodeLoginRateLimitExceededError,
             )
 
             raise EmailCodeLoginRateLimitExceededError()
 
-        code = "".join([str(random.randint(0, 9)) for _ in range(6)])
+        # if debug mode, force set code to 111111
+        if DeploymentConfig().DEBUG:
+            code = "111111"  # TODO: ytqh move this to config
+        else:
+            code = "".join([str(random.randint(0, 9)) for _ in range(6)])
+
         token = TokenManager.generate_token(
             account=account,
             email=email,