tjdify/api/services/account_service.py

# -*- 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')