fix: get_admin_through_phone check logic

pull/21891/head
ytqh 1 year ago
parent d11ebef66f
commit c90590ed62

@ -8,11 +8,6 @@ from datetime import UTC, datetime, timedelta
from hashlib import sha256 from hashlib import sha256
from typing import Any, Optional, cast from typing import Any, Optional, cast
from pydantic import BaseModel # type: ignore
from sqlalchemy import func # type: ignore
from sqlalchemy.orm import Session
from werkzeug.exceptions import Unauthorized
from configs import dify_config from configs import dify_config
from constants.languages import language_timezone_mapping, languages from constants.languages import language_timezone_mapping, languages
from events.tenant_event import tenant_was_created from events.tenant_event import tenant_was_created
@ -22,42 +17,36 @@ from libs.helper import RateLimiter, TokenManager
from libs.passport import PassportService from libs.passport import PassportService
from libs.password import compare_password, hash_password, valid_password from libs.password import compare_password, hash_password, valid_password
from libs.rsa import generate_key_pair from libs.rsa import generate_key_pair
from models.account import ( from models.account import (Account, AccountIntegrate, AccountStatus, Tenant,
Account, TenantAccountJoin, TenantAccountJoinRole,
AccountIntegrate, TenantAccountRole, TenantStatus)
AccountStatus,
Tenant,
TenantAccountJoin,
TenantAccountJoinRole,
TenantAccountRole,
TenantStatus,
)
from models.model import DifySetup from models.model import DifySetup
from pydantic import BaseModel # type: ignore
from services.billing_service import BillingService from services.billing_service import BillingService
from services.errors.account import ( from services.errors.account import (AccountAlreadyInTenantError,
AccountAlreadyInTenantError, AccountLoginError, AccountNotFoundError,
AccountLoginError, AccountNotLinkTenantError,
AccountNotFoundError, AccountPasswordError,
AccountNotLinkTenantError, AccountRegisterError,
AccountPasswordError, CannotOperateSelfError,
AccountRegisterError, CurrentPasswordIncorrectError,
CannotOperateSelfError, InvalidActionError,
CurrentPasswordIncorrectError, LinkAccountIntegrateError,
InvalidActionError, MemberNotInTenantError, NoPermissionError,
LinkAccountIntegrateError, RoleAlreadyAssignedError,
MemberNotInTenantError, TenantNotFoundError)
NoPermissionError,
RoleAlreadyAssignedError,
TenantNotFoundError,
)
from services.errors.workspace import WorkSpaceNotAllowedCreateError from services.errors.workspace import WorkSpaceNotAllowedCreateError
from services.feature_service import FeatureService from services.feature_service import FeatureService
from sqlalchemy import func # type: ignore
from sqlalchemy.orm import Session
from tasks.delete_account_task import delete_account_task from tasks.delete_account_task import delete_account_task
from tasks.mail_account_deletion_task import send_account_deletion_verification_code 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 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_invite_member_task import send_invite_member_mail_task
from tasks.mail_reset_password_task import send_reset_password_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task
from tasks.phone_sms_code_login import send_phone_sms_code_login_task from tasks.phone_sms_code_login import send_phone_sms_code_login_task
from werkzeug.exceptions import Unauthorized
class TokenPair(BaseModel): class TokenPair(BaseModel):
@ -306,7 +295,8 @@ class AccountService:
def send_account_deletion_verification_email(cls, account: Account, code: str): def send_account_deletion_verification_email(cls, account: Account, code: str):
email = account.email email = account.email
if cls.email_code_account_deletion_rate_limiter.is_rate_limited(email): if cls.email_code_account_deletion_rate_limiter.is_rate_limited(email):
from controllers.console.auth.error import EmailCodeAccountDeletionRateLimitExceededError from controllers.console.auth.error import \
EmailCodeAccountDeletionRateLimitExceededError
raise EmailCodeAccountDeletionRateLimitExceededError() raise EmailCodeAccountDeletionRateLimitExceededError()
@ -444,7 +434,8 @@ class AccountService:
raise ValueError("Email must be provided.") raise ValueError("Email must be provided.")
if cls.reset_password_rate_limiter.is_rate_limited(account_email): if cls.reset_password_rate_limiter.is_rate_limited(account_email):
from controllers.console.auth.error import PasswordResetRateLimitExceededError from controllers.console.auth.error import \
PasswordResetRateLimitExceededError
raise PasswordResetRateLimitExceededError() raise PasswordResetRateLimitExceededError()
@ -496,7 +487,8 @@ class AccountService:
if dify_config.DEBUG_CODE_FOR_LOGIN: if dify_config.DEBUG_CODE_FOR_LOGIN:
code = dify_config.DEBUG_CODE_FOR_LOGIN code = dify_config.DEBUG_CODE_FOR_LOGIN
elif cls.email_code_login_rate_limiter.is_rate_limited(email): elif cls.email_code_login_rate_limiter.is_rate_limited(email):
from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError from controllers.console.auth.error import \
EmailCodeLoginRateLimitExceededError
raise EmailCodeLoginRateLimitExceededError() raise EmailCodeLoginRateLimitExceededError()
else: else:
@ -640,18 +632,16 @@ class AccountService:
""" """
Get admin account through phone number. Get admin account through phone number.
Args:
phone: The phone number to search for
Returns None if no admin account with this phone number exists. Returns None if no admin account with this phone number exists.
Raises AccountRegisterError if account is in freeze status. Raises Unauthorized if account is banned.
""" """
# We'll check if account is banned or frozen using the phone number # Query directly with phone number first
# This implementation assumes phone numbers are unique like emails
# Find account with end_admin role by phone number
admin_account = ( admin_account = (
db.session.query(Account) db.session.query(Account)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(Account.phone == phone) .filter(Account.phone == phone)
.filter(TenantAccountJoin.role == TenantAccountJoinRole.END_ADMIN.value)
.first() .first()
) )
@ -661,6 +651,29 @@ class AccountService:
if admin_account.status == AccountStatus.BANNED.value: if admin_account.status == AccountStatus.BANNED.value:
raise Unauthorized("Account is banned.") raise Unauthorized("Account is banned.")
organization_id = admin_account.current_organization_id
if not organization_id:
logging.warning(f"Account {admin_account.id} is not a member of any organization.")
return None
# If organization_id is provided, check if account is an admin member of that organization
from models.organization import OrganizationMember, OrganizationRole
org_member = (
db.session.query(OrganizationMember)
.filter(
OrganizationMember.organization_id == organization_id,
OrganizationMember.account_id == admin_account.id,
OrganizationMember.role == OrganizationRole.ADMIN
)
.first()
)
if not org_member:
logging.warning(f"Account {admin_account.id} is not a member of any organization.")
return None
return admin_account return admin_account
@classmethod @classmethod

Loading…
Cancel
Save