# -*- coding:utf-8 -*- import base64 import logging import secrets import uuid from datetime import datetime from hashlib import sha256 from typing import Optional from flask import session from sqlalchemy import func from events.tenant_event import tenant_was_created from extensions.ext_redis import redis_client from services.errors.account import AccountLoginError, CurrentPasswordIncorrectError, LinkAccountIntegrateError, \ TenantNotFound, AccountNotLinkTenantError, InvalidActionError, CannotOperateSelfError, MemberNotInTenantError, \ RoleAlreadyAssignedError, NoPermissionError, AccountRegisterError, AccountAlreadyInTenantError from libs.helper import get_remote_ip from libs.password import compare_password, hash_password from libs.rsa import generate_key_pair from models.account import * from tasks.mail_invite_member_task import send_invite_member_mail_task class AccountService: @staticmethod def load_user(account_id: int) -> Account: # todo: used by flask_login pass @staticmethod def authenticate(email: str, password: str) -> Account: """authenticate account with email and password""" account = Account.query.filter_by(email=email).first() if not account: raise AccountLoginError('Invalid email or password.') if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value: raise AccountLoginError('Account is banned or closed.') if account.status == AccountStatus.PENDING.value: account.status = AccountStatus.ACTIVE.value account.initialized_at = datetime.utcnow() db.session.commit() if account.password is None or not compare_password(password, account.password, account.password_salt): raise AccountLoginError('Invalid email or password.') return account @staticmethod def update_account_password(account, password, new_password): """update account password""" if account.password and not compare_password(password, account.password, account.password_salt): raise CurrentPasswordIncorrectError("Current password is incorrect.") # generate password salt salt = secrets.token_bytes(16) base64_salt = base64.b64encode(salt).decode() # encrypt password with salt password_hashed = hash_password(new_password, salt) base64_password_hashed = base64.b64encode(password_hashed).decode() account.password = base64_password_hashed account.password_salt = base64_salt db.session.commit() return account @staticmethod def create_account(email: str, name: str, password: str = None, interface_language: str = 'en-US', interface_theme: str = 'light', timezone: str = 'America/New_York', ) -> Account: """create account""" account = Account() account.email = email account.name = name if password: # generate password salt salt = secrets.token_bytes(16) base64_salt = base64.b64encode(salt).decode() # encrypt password with salt password_hashed = hash_password(password, salt) base64_password_hashed = base64.b64encode(password_hashed).decode() account.password = base64_password_hashed account.password_salt = base64_salt account.interface_language = interface_language account.interface_theme = interface_theme if interface_language == 'zh-Hans': account.timezone = 'Asia/Shanghai' else: account.timezone = timezone db.session.add(account) db.session.commit() return account @staticmethod def link_account_integrate(provider: str, open_id: str, account: Account) -> None: """Link account integrate""" try: # Query whether there is an existing binding record for the same provider account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(account_id=account.id, provider=provider).first() if account_integrate: # If it exists, update the record account_integrate.open_id = open_id account_integrate.encrypted_token = "" # todo account_integrate.updated_at = datetime.utcnow() else: # If it does not exist, create a new record account_integrate = AccountIntegrate(account_id=account.id, provider=provider, open_id=open_id, encrypted_token="") db.session.add(account_integrate) db.session.commit() logging.info(f'Account {account.id} linked {provider} account {open_id}.') except Exception as e: logging.exception(f'Failed to link {provider} account {open_id} to Account {account.id}') raise LinkAccountIntegrateError('Failed to link account.') from e @staticmethod def close_account(account: Account) -> None: """todo: Close account""" account.status = AccountStatus.CLOSED.value db.session.commit() @staticmethod def update_account(account, **kwargs): """Update account fields""" for field, value in kwargs.items(): if hasattr(account, field): setattr(account, field, value) else: raise AttributeError(f"Invalid field: {field}") db.session.commit() return account @staticmethod def update_last_login(account: Account, request) -> None: """Update last login time and ip""" account.last_login_at = datetime.utcnow() account.last_login_ip = get_remote_ip(request) db.session.add(account) db.session.commit() logging.info(f'Account {account.id} logged in successfully.') class TenantService: @staticmethod def create_tenant(name: str) -> Tenant: """Create tenant""" tenant = Tenant(name=name) db.session.add(tenant) db.session.commit() tenant.encrypt_public_key = generate_key_pair(tenant.id) db.session.commit() return tenant @staticmethod def create_tenant_member(tenant: Tenant, account: Account, role: str = 'normal') -> TenantAccountJoin: """Create tenant member""" if role == TenantAccountJoinRole.OWNER.value: if TenantService.has_roles(tenant, [TenantAccountJoinRole.OWNER]): 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) db.session.commit() return ta @staticmethod def get_join_tenants(account: Account) -> List[Tenant]: """Get account join tenants""" return db.session.query(Tenant).join( TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id ).filter(TenantAccountJoin.account_id == account.id).all() @staticmethod def get_current_tenant_by_account(account: Account): """Get tenant by account and add the role""" tenant = account.current_tenant if not tenant: raise TenantNotFound("Tenant not found.") ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first() if ta: tenant.role = ta.role else: raise TenantNotFound("Tenant not found for the account.") return tenant @staticmethod def switch_tenant(account: Account, tenant_id: int = None) -> None: """Switch the current workspace for the account""" if not tenant_id: tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id).first() else: tenant_account_join = TenantAccountJoin.query.filter_by(account_id=account.id, tenant_id=tenant_id).first() # Check if the tenant exists and the account is a member of the tenant if not tenant_account_join: raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.") # Set the current tenant for the account account.current_tenant_id = tenant_account_join.tenant_id session['workspace_id'] = account.current_tenant.id @staticmethod def get_tenant_members(tenant: Tenant) -> List[Account]: """Get tenant members""" query = ( db.session.query(Account, TenantAccountJoin.role) .select_from(Account) .join( TenantAccountJoin, Account.id == TenantAccountJoin.account_id ) .filter(TenantAccountJoin.tenant_id == tenant.id) ) # Initialize an empty list to store the updated accounts updated_accounts = [] for account, role in query: account.role = role updated_accounts.append(account) return updated_accounts @staticmethod def has_roles(tenant: Tenant, roles: List[TenantAccountJoinRole]) -> bool: """Check if user has any of the given roles for a tenant""" if not all(isinstance(role, TenantAccountJoinRole) for role in roles): raise ValueError('all roles must be TenantAccountJoinRole') return db.session.query(TenantAccountJoin).filter( TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles]) ).first() is not None @staticmethod def get_user_role(account: Account, tenant: Tenant) -> Optional[TenantAccountJoinRole]: """Get the role of the current account for a given tenant""" join = db.session.query(TenantAccountJoin).filter( TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id ).first() return join.role if join else None @staticmethod def get_tenant_count() -> int: """Get tenant count""" return db.session.query(func.count(Tenant.id)).scalar() @staticmethod def check_member_permission(tenant: Tenant, operator: Account, member: Account, action: str) -> None: """Check member permission""" perms = { 'add': ['owner', 'admin'], 'remove': ['owner'], 'update': ['owner'] } if action not in ['add', 'remove', 'update']: raise InvalidActionError("Invalid action.") if member: if operator.id == member.id: raise CannotOperateSelfError("Cannot operate self.") ta_operator = TenantAccountJoin.query.filter_by( tenant_id=tenant.id, account_id=operator.id ).first() if not ta_operator or ta_operator.role not in perms[action]: raise NoPermissionError(f'No permission to {action} member.') @staticmethod def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account) -> None: """Remove member from tenant""" if operator.id == account.id and TenantService.check_member_permission(tenant, operator, account, 'remove'): raise CannotOperateSelfError("Cannot operate self.") ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first() if not ta: raise MemberNotInTenantError("Member not in tenant.") db.session.delete(ta) account.initialized_at = None account.status = AccountStatus.PENDING.value account.password = None account.password_salt = None db.session.commit() @staticmethod def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account) -> None: """Update member role""" TenantService.check_member_permission(tenant, operator, member, 'update') target_member_join = TenantAccountJoin.query.filter_by( tenant_id=tenant.id, account_id=member.id ).first() if target_member_join.role == new_role: raise RoleAlreadyAssignedError("The provided role is already assigned to the member.") if new_role == 'owner': # Find the current owner and change their role to 'admin' current_owner_join = TenantAccountJoin.query.filter_by( tenant_id=tenant.id, role='owner' ).first() current_owner_join.role = 'admin' # Update the role of the target member target_member_join.role = new_role db.session.commit() @staticmethod def dissolve_tenant(tenant: Tenant, operator: Account) -> None: """Dissolve tenant""" if not TenantService.check_member_permission(tenant, operator, operator, 'remove'): raise NoPermissionError('No permission to dissolve tenant.') db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id).delete() db.session.delete(tenant) db.session.commit() class RegisterService: @classmethod def register(cls, email, name, password: str = None, open_id: str = None, provider: str = None) -> Account: db.session.begin_nested() """Register account""" try: account = AccountService.create_account(email, name, password) account.status = AccountStatus.ACTIVE.value account.initialized_at = datetime.utcnow() if open_id is not None or provider is not None: AccountService.link_account_integrate(provider, open_id, account) tenant = TenantService.create_tenant(f"{account.name}'s Workspace") TenantService.create_tenant_member(tenant, account, role='owner') account.current_tenant = tenant db.session.commit() except Exception as e: db.session.rollback() # todo: do not work logging.error(f'Register failed: {e}') raise AccountRegisterError(f'Registration failed: {e}') from e tenant_was_created.send(tenant) return account @classmethod def invite_new_member(cls, tenant: Tenant, email: str, role: str = 'normal', inviter: Account = None) -> str: """Invite new member""" account = Account.query.filter_by(email=email).first() if not account: TenantService.check_member_permission(tenant, inviter, None, 'add') name = email.split('@')[0] account = AccountService.create_account(email, name) account.status = AccountStatus.PENDING.value db.session.commit() else: TenantService.check_member_permission(tenant, inviter, account, 'add') ta = TenantAccountJoin.query.filter_by( tenant_id=tenant.id, account_id=account.id ).first() if ta: raise AccountAlreadyInTenantError("Account already in tenant.") TenantService.create_tenant_member(tenant, account, role) token = cls.generate_invite_token(tenant, account) # send email send_invite_member_mail_task.delay( to=email, token=cls.generate_invite_token(tenant, account), inviter_name=inviter.name if inviter else 'Dify', workspace_id=tenant.id, workspace_name=tenant.name, ) return token @classmethod def generate_invite_token(cls, tenant: Tenant, account: Account) -> str: token = str(uuid.uuid4()) email_hash = sha256(account.email.encode()).hexdigest() cache_key = 'member_invite_token:{}, {}:{}'.format(str(tenant.id), email_hash, token) redis_client.setex(cache_key, 3600, str(account.id)) return token @classmethod def revoke_token(cls, workspace_id: str, email: str, token: str): email_hash = sha256(email.encode()).hexdigest() cache_key = 'member_invite_token:{}, {}:{}'.format(workspace_id, email_hash, token) redis_client.delete(cache_key) @classmethod def get_account_if_token_valid(cls, workspace_id: str, email: str, token: str) -> Optional[Account]: tenant = db.session.query(Tenant).filter( Tenant.id == workspace_id, Tenant.status == 'normal' ).first() if not tenant: return None tenant_account = db.session.query(Account, TenantAccountJoin.role).join( TenantAccountJoin, Account.id == TenantAccountJoin.account_id ).filter(Account.email == email, TenantAccountJoin.tenant_id == tenant.id).first() if not tenant_account: return None account_id = cls._get_account_id_by_invite_token(workspace_id, email, token) if not account_id: return None account = tenant_account[0] if not account: return None if account_id != str(account.id): return None return account @classmethod def _get_account_id_by_invite_token(cls, workspace_id: str, email: str, token: str) -> Optional[str]: email_hash = sha256(email.encode()).hexdigest() cache_key = 'member_invite_token:{}, {}:{}'.format(workspace_id, email_hash, token) account_id = redis_client.get(cache_key) if not account_id: return None return account_id.decode('utf-8')