From 22bdfb7e56019452dcb7e95725ee550b17b7a8df Mon Sep 17 00:00:00 2001
From: Joe <79627742+ZhouhaoJiang@users.noreply.github.com>
Date: Wed, 23 Oct 2024 10:59:30 +0800
Subject: [PATCH] Feat/optimize login (#9642)

---
 api/controllers/console/auth/oauth.py |  8 +++-----
 api/services/account_service.py       | 12 ++++++------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py
index 282e69448e..d27e3353c9 100644
--- a/api/controllers/console/auth/oauth.py
+++ b/api/controllers/console/auth/oauth.py
@@ -94,17 +94,15 @@ class OAuthCallback(Resource):
             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:
+        except (WorkSpaceNotFoundError, 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
+        if account.status == AccountStatus.BANNED.value:
+            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Account is banned.")
 
         if account.status == AccountStatus.PENDING.value:
             account.status = AccountStatus.ACTIVE.value
diff --git a/api/services/account_service.py b/api/services/account_service.py
index 529b716773..412685147c 100644
--- a/api/services/account_service.py
+++ b/api/services/account_service.py
@@ -98,8 +98,8 @@ class AccountService:
         if not account:
             return None
 
-        if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}:
-            raise Unauthorized("Account is banned or closed.")
+        if account.status == AccountStatus.BANNED.value:
+            raise Unauthorized("Account is banned.")
 
         current_tenant = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
         if current_tenant:
@@ -143,8 +143,8 @@ class AccountService:
         if not account:
             raise AccountNotFoundError()
 
-        if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}:
-            raise AccountLoginError("Account is banned or closed.")
+        if account.status == AccountStatus.BANNED.value:
+            raise AccountLoginError("Account is banned.")
 
         if password and invite_token and account.password is None:
             # if invite_token is valid, set password and password_salt
@@ -408,8 +408,8 @@ class AccountService:
         if not account:
             return None
 
-        if account.status in {AccountStatus.BANNED.value, AccountStatus.CLOSED.value}:
-            raise Unauthorized("Account is banned or closed.")
+        if account.status == AccountStatus.BANNED.value:
+            raise Unauthorized("Account is banned.")
 
         return account