tjdify/api/services/account_service.py

# -*- coding:utf-8 -*-
import base64
import json
import logging
import secrets
import uuid
from datetime import datetime, timedelta
from hashlib import sha256
from typing import Optional

from werkzeug.exceptions import Forbidden, Unauthorized
from flask import session, current_app
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 libs.passport import PassportService
from models.account import *
from tasks.mail_invite_member_task import send_invite_member_mail_task

def _create_tenant_for_account(account):
    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")

    TenantService.create_tenant_member(tenant, account, role='owner')
    account.current_tenant = tenant

    return tenant


class AccountService:

    @staticmethod
    def load_user(user_id: str) -> Account:
        # todo: used by flask_login
        if '.' in user_id:
            tenant_id, account_id = user_id.split('.')
        else:
            account_id = user_id

        account = db.session.query(Account).filter(Account.id == account_id).first()

        if account:
            if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
                raise Forbidden('Account is banned or closed.')

            workspace_id = session.get('workspace_id')
            if workspace_id:
                tenant_account_join = db.session.query(TenantAccountJoin).filter(
                    TenantAccountJoin.account_id == account.id,
                    TenantAccountJoin.tenant_id == workspace_id
                ).first()

                if not tenant_account_join:
                    tenant_account_join = db.session.query(TenantAccountJoin).filter(
                        TenantAccountJoin.account_id == account.id).first()

                    if tenant_account_join:
                        account.current_tenant_id = tenant_account_join.tenant_id
                    else:
                        _create_tenant_for_account(account)
                    session['workspace_id'] = account.current_tenant_id
                else:
                    account.current_tenant_id = workspace_id
            else:
                tenant_account_join = db.session.query(TenantAccountJoin).filter(
                    TenantAccountJoin.account_id == account.id).first()
                if tenant_account_join:
                    account.current_tenant_id = tenant_account_join.tenant_id
                else:
                    _create_tenant_for_account(account)
                session['workspace_id'] = account.current_tenant_id

            current_time = datetime.utcnow()

            # update last_active_at when last_active_at is more than 10 minutes ago
            if current_time - account.last_active_at > timedelta(minutes=10):
                account.last_active_at = current_time
                db.session.commit()

        return account
    
    @staticmethod
    def get_account_jwt_token(account):
        payload = {
            "user_id": account.id,
            "exp": datetime.utcnow() + timedelta(days=30),
            "iss":  current_app.config['EDITION'],
            "sub": 'Console API Passport',
        }

        token = PassportService().issue(payload)
        return token

    @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 _get_invitation_token_key(cls, token: str) -> str:
        return f'member_invite:token:{token}'

    @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=token,
            inviter_name=inviter.name if inviter else 'Dify',
            workspace_name=tenant.name,
        )

        return token

    @classmethod
    def generate_invite_token(cls, tenant: Tenant, account: Account) -> str:
        token = str(uuid.uuid4())
        invitation_data = {
            'account_id': account.id,
            'email': account.email,
            'workspace_id': tenant.id,
        }
        redis_client.setex(
            cls._get_invitation_token_key(token),
            3600,
            json.dumps(invitation_data)
        )
        return token

    @classmethod
    def revoke_token(cls, workspace_id: str, email: str, token: str):
        if workspace_id and email:
            email_hash = sha256(email.encode()).hexdigest()
            cache_key = 'member_invite_token:{}, {}:{}'.format(workspace_id, email_hash, token)
            redis_client.delete(cache_key)
        else:
            redis_client.delete(cls._get_invitation_token_key(token))

    @classmethod
    def get_invitation_if_token_valid(cls, workspace_id: str, email: str, token: str) -> Optional[Account]:
        invitation_data = cls._get_invitation_by_token(token, workspace_id, email)
        if not invitation_data:
            return None

        tenant = db.session.query(Tenant).filter(
            Tenant.id == invitation_data['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 == invitation_data['email'], TenantAccountJoin.tenant_id == tenant.id).first()

        if not tenant_account:
            return None

        account = tenant_account[0]
        if not account:
            return None

        if invitation_data['account_id'] != str(account.id):
            return None

        return {
                'account': account,
                'data': invitation_data,
                'tenant': tenant,
                }

    @classmethod
    def _get_invitation_by_token(cls, token: str, workspace_id: str, email: str) -> Optional[str]:
        if workspace_id is not None and email is not None:
            email_hash = sha256(email.encode()).hexdigest()
            cache_key = f'member_invite_token:{workspace_id}, {email_hash}:{token}'
            account_id = redis_client.get(cache_key)

            if not account_id:
                return None

            return {
                'account_id': account_id.decode('utf-8'),
                'email': email,
                'workspace_id': workspace_id,
            }
        else:
            data = redis_client.get(cls._get_invitation_token_key(token))
            if not data:
                return None

            invitation = json.loads(data)
            return invitation