@ -1,6 +1,7 @@
import base64
import json
import logging
import random
import secrets
import uuid
from datetime import datetime , timedelta , timezone
@ -34,7 +35,9 @@ from models.model import DifySetup
from services . errors . account import (
AccountAlreadyInTenantError ,
AccountLoginError ,
AccountNotFoundError ,
AccountNotLinkTenantError ,
AccountPasswordError ,
AccountRegisterError ,
CannotOperateSelfError ,
CurrentPasswordIncorrectError ,
@ -42,10 +45,12 @@ from services.errors.account import (
LinkAccountIntegrateError ,
MemberNotInTenantError ,
NoPermissionError ,
RateLimitExceededError ,
RoleAlreadyAssignedError ,
TenantNotFoundError ,
)
from services . errors . workspace import WorkSpaceNotAllowedCreateError
from services . feature_service import FeatureService
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_reset_password_task import send_reset_password_mail_task
@ -61,7 +66,11 @@ REFRESH_TOKEN_EXPIRY = timedelta(days=30)
class AccountService :
reset_password_rate_limiter = RateLimiter ( prefix = " reset_password_rate_limit " , max_attempts = 5 , time_window = 60 * 60 )
reset_password_rate_limiter = RateLimiter ( prefix = " reset_password_rate_limit " , max_attempts = 1 , time_window = 60 * 1 )
email_code_login_rate_limiter = RateLimiter (
prefix = " email_code_login_rate_limit " , max_attempts = 1 , time_window = 60 * 1
)
LOGIN_MAX_ERROR_LIMITS = 5
@staticmethod
def _get_refresh_token_key ( refresh_token : str ) - > str :
@ -127,23 +136,34 @@ class AccountService:
return token
@staticmethod
def authenticate ( email : str , password : str - > Account :
def authenticate ( email : str , password : str , invite_token : Optional [ str ] = None ) - > Account :
""" authenticate account with email and password """
account = Account . query . filter_by ( email = email ) . first ( )
if not account :
raise Account LoginError( " Invalid email or password. " )
raise Account NotFoundError( )
if account . status in { AccountStatus . BANNED . value , AccountStatus . CLOSED . value } :
raise AccountLoginError ( " Account is banned or closed. " )
if password and invite_token and account . password is None :
# if invite_token is valid, set password and password_salt
salt = secrets . token_bytes ( 16 )
base64_salt = base64 . b64encode ( salt ) . decode ( )
password_hashed = hash_password ( password , salt )
base64_password_hashed = base64 . b64encode ( password_hashed ) . decode ( )
account . password = base64_password_hashed
account . password_salt = base64_salt
if account . password is None or not compare_password ( password , account . password , account . password_salt ) :
raise AccountPasswordError ( " Invalid email or password. " )
if account . status == AccountStatus . PENDING . value :
account . status = AccountStatus . ACTIVE . value
account . initialized_at = datetime . now ( timezone . utc ) . replace ( tzinfo = None )
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
@ -169,9 +189,18 @@ class AccountService:
@staticmethod
def create_account (
email : str , name : str , interface_language : str , password : Optional [ str ] = None , interface_theme : str = " light "
email : str ,
name : str ,
interface_language : str ,
password : Optional [ str ] = None ,
interface_theme : str = " light " ,
is_setup : Optional [ bool ] = False ,
) - > Account :
""" create account """
if not FeatureService . get_system_features ( ) . is_allow_register and not is_setup :
from controllers . console . error import NotAllowedRegister
raise NotAllowedRegister ( )
account = Account ( )
account . email = email
account . name = name
@ -198,6 +227,19 @@ class AccountService:
db . session . commit ( )
return account
@staticmethod
def create_account_and_tenant (
email : str , name : str , interface_language : str , password : Optional [ str ] = None
) - > Account :
""" create account """
account = AccountService . create_account (
email = email , name = name , interface_language = interface_language , password = password
)
TenantService . create_owner_tenant_if_not_exist ( account = account )
return account
@staticmethod
def link_account_integrate ( provider : str , open_id : str , account : Account ) - > None :
""" Link account integrate """
@ -256,6 +298,10 @@ class AccountService:
if ip_address :
AccountService . update_login_info ( account = account , ip_address = ip_address )
if account . status == AccountStatus . PENDING . value :
account . status = AccountStatus . ACTIVE . value
db . session . commit ( )
access_token = AccountService . get_account_jwt_token ( account = account )
refresh_token = _generate_refresh_token ( )
@ -294,13 +340,29 @@ class AccountService:
return AccountService . load_user ( account_id )
@classmethod
def send_reset_password_email ( cls , account ) :
if cls . reset_password_rate_limiter . is_rate_limited ( account . email ) :
raise RateLimitExceededError ( f " Rate limit exceeded for email: { account . email } . Please try again later. " )
def send_reset_password_email (
cls ,
account : Optional [ Account ] = None ,
email : Optional [ str ] = None ,
language : Optional [ str ] = " en-US " ,
) :
account_email = account . email if account else email
token = TokenManager . generate_token ( account , " reset_password " )
send_reset_password_mail_task . delay ( language = account . interface_language , to = account . email , token = token )
cls . reset_password_rate_limiter . increment_rate_limit ( account . email )
if cls . reset_password_rate_limiter . is_rate_limited ( account_email ) :
from controllers . console . auth . error import PasswordResetRateLimitExceededError
raise PasswordResetRateLimitExceededError ( )
code = " " . join ( [ str ( random . randint ( 0 , 9 ) ) for _ in range ( 6 ) ] )
token = TokenManager . generate_token (
account = account , email = email , token_type = " reset_password " , additional_data = { " code " : code }
)
send_reset_password_mail_task . delay (
language = language ,
to = account_email ,
code = code ,
)
cls . reset_password_rate_limiter . increment_rate_limit ( account_email )
return token
@classmethod
@ -311,11 +373,125 @@ class AccountService:
def get_reset_password_data ( cls , token : str ) - > Optional [ dict [ str , Any ] ] :
return TokenManager . get_token_data ( token , " reset_password " )
@classmethod
def send_email_code_login_email (
cls , account : Optional [ Account ] = None , email : Optional [ str ] = None , language : Optional [ str ] = " en-US "
) :
if cls . email_code_login_rate_limiter . is_rate_limited ( email ) :
from controllers . console . auth . error import EmailCodeLoginRateLimitExceededError
raise EmailCodeLoginRateLimitExceededError ( )
code = " " . join ( [ str ( random . randint ( 0 , 9 ) ) for _ in range ( 6 ) ] )
token = TokenManager . generate_token (
account = account , email = email , token_type = " email_code_login " , additional_data = { " code " : code }
)
send_email_code_login_mail_task . delay (
language = language ,
to = account . email if account else email ,
code = code ,
)
cls . email_code_login_rate_limiter . increment_rate_limit ( email )
return token
@classmethod
def get_email_code_login_data ( cls , token : str ) - > Optional [ dict [ str , Any ] ] :
return TokenManager . get_token_data ( token , " email_code_login " )
@classmethod
def revoke_email_code_login_token ( cls , token : str ) :
TokenManager . revoke_token ( token , " email_code_login " )
@classmethod
def get_user_through_email ( cls , email : str ) :
account = db . session . query ( Account ) . filter ( Account . email == email ) . first ( )
if not account :
return None
if account . status in { AccountStatus . BANNED . value , AccountStatus . CLOSED . value } :
raise Unauthorized ( " Account is banned or closed. " )
return account
@staticmethod
def add_login_error_rate_limit ( email : str ) - > None :
key = f " login_error_rate_limit: { email } "
count = redis_client . get ( key )
if count is None :
count = 0
count = int ( count ) + 1
redis_client . setex ( key , 60 * 60 * 24 , count )
@staticmethod
def is_login_error_rate_limit ( email : str ) - > bool :
key = f " login_error_rate_limit: { email } "
count = redis_client . get ( key )
if count is None :
return False
count = int ( count )
if count > AccountService . LOGIN_MAX_ERROR_LIMITS :
return True
return False
@staticmethod
def reset_login_error_rate_limit ( email : str ) :
key = f " login_error_rate_limit: { email } "
redis_client . delete ( key )
@staticmethod
def is_email_send_ip_limit ( ip_address : str ) :
minute_key = f " email_send_ip_limit_minute: { ip_address } "
freeze_key = f " email_send_ip_limit_freeze: { ip_address } "
hour_limit_key = f " email_send_ip_limit_hour: { ip_address } "
# check ip is frozen
if redis_client . get ( freeze_key ) :
return True
# check current minute count
current_minute_count = redis_client . get ( minute_key )
if current_minute_count is None :
current_minute_count = 0
current_minute_count = int ( current_minute_count )
# check current hour count
if current_minute_count > dify_config . EMAIL_SEND_IP_LIMIT_PER_MINUTE :
hour_limit_count = redis_client . get ( hour_limit_key )
if hour_limit_count is None :
hour_limit_count = 0
hour_limit_count = int ( hour_limit_count )
if hour_limit_count > = 1 :
redis_client . setex ( freeze_key , 60 * 60 , 1 )
return True
else :
redis_client . setex ( hour_limit_key , 60 * 10 , hour_limit_count + 1 ) # first time limit 10 minutes
# add hour limit count
redis_client . incr ( hour_limit_key )
redis_client . expire ( hour_limit_key , 60 * 60 )
return True
redis_client . setex ( minute_key , 60 , current_minute_count + 1 )
redis_client . expire ( minute_key , 60 )
return False
def _get_login_cache_key ( * , account_id : str , token : str ) :
return f " account_login: { account_id } : { token } "
class TenantService :
@staticmethod
def create_tenant ( name : str ) - > Tenant :
def create_tenant ( name : str , is_setup : Optional [ bool ] = False ) - > Tenant :
""" Create tenant """
if not FeatureService . get_system_features ( ) . is_allow_create_workspace and not is_setup :
from controllers . console . error import NotAllowedCreateWorkspace
raise NotAllowedCreateWorkspace ( )
tenant = Tenant ( name = name )
db . session . add ( tenant )
@ -326,8 +502,12 @@ class TenantService:
return tenant
@staticmethod
def create_owner_tenant_if_not_exist ( account : Account , name : Optional [ str ] = None ) :
def create_owner_tenant_if_not_exist (
account : Account , name : Optional [ str ] = None , is_setup : Optional [ bool ] = False
) :
""" Create owner tenant if not exist """
if not FeatureService . get_system_features ( ) . is_allow_create_workspace and not is_setup :
raise WorkSpaceNotAllowedCreateError ( )
available_ta = (
TenantAccountJoin . query . filter_by ( account_id = account . id ) . order_by ( TenantAccountJoin . id . asc ( ) ) . first ( )
)
@ -336,9 +516,9 @@ class TenantService:
return
if name :
tenant = TenantService . create_tenant ( name )
tenant = TenantService . create_tenant ( name = name , is_setup = is_setup )
else :
tenant = TenantService . create_tenant ( f " { account . name } ' s Workspace " )
tenant = TenantService . create_tenant ( name = f " { account . name } ' s Workspace " , is_setup = is_setup )
TenantService . create_tenant_member ( tenant , account , role = " owner " )
account . current_tenant = tenant
db . session . commit ( )
@ -352,8 +532,13 @@ class TenantService:
logging . error ( f " Tenant { tenant . id } has already an owner. " )
raise Exception ( " Tenant already has an owner. " )
ta = db . session . query ( TenantAccountJoin ) . filter_by ( tenant_id = tenant . id , account_id = account . id ) . first ( )
if ta :
ta . role = role
else :
ta = TenantAccountJoin ( tenant_id = tenant . id , account_id = account . id , role = role )
db . session . add ( ta )
db . session . commit ( )
return ta
@ -570,12 +755,13 @@ class RegisterService:
name = name ,
interface_language = languages [ 0 ] ,
password = password ,
is_setup = True ,
)
account . last_login_ip = ip_address
account . initialized_at = datetime . now ( timezone . utc ) . replace ( tzinfo = None )
TenantService . create_owner_tenant_if_not_exist ( account )
TenantService . create_owner_tenant_if_not_exist ( account = account , is_setup = True )
dify_setup = DifySetup ( version = dify_config . CURRENT_VERSION )
db . session . add ( dify_setup )
@ -600,27 +786,33 @@ class RegisterService:
provider : Optional [ str ] = None ,
language : Optional [ str ] = None ,
status : Optional [ AccountStatus ] = None ,
is_setup : Optional [ bool ] = False ,
) - > Account :
db . session . begin_nested ( )
""" Register account """
try :
account = AccountService . create_account (
email = email , name = name , interface_language = language or languages [ 0 ] , password = password
email = email ,
name = name ,
interface_language = language or languages [ 0 ] ,
password = password ,
is_setup = is_setup ,
)
account . status = AccountStatus . ACTIVE . value if not status else status . value
account . initialized_at = datetime . now ( timezone . utc ) . replace ( tzinfo = None )
if open_id is not None or provider is not None :
AccountService . link_account_integrate ( provider , open_id , account )
if dify_config . EDITION != " SELF_HOSTED " :
tenant = TenantService . create_tenant ( f " { account . name } ' s Workspace " )
if FeatureService . get_system_features ( ) . is_allow_create_workspace :
tenant = TenantService . create_tenant ( f " { account . name } ' s Workspace " )
TenantService . create_tenant_member ( tenant , account , role = " owner " )
account . current_tenant = tenant
tenant_was_created . send ( tenant )
db . session . commit ( )
except WorkSpaceNotAllowedCreateError :
db . session . rollback ( )
except Exception as e :
db . session . rollback ( )
logging . error ( f " Register failed: { e } " )
@ -639,7 +831,9 @@ class RegisterService:
TenantService . check_member_permission ( tenant , inviter , None , " add " )
name = email . split ( " @ " ) [ 0 ]
account = cls . register ( email = email , name = name , language = language , status = AccountStatus . PENDING )
account = cls . register (
email = email , name = name , language = language , status = AccountStatus . PENDING , is_setup = True
)
# Create new tenant member for invited tenant
TenantService . create_tenant_member ( tenant , account , role )
TenantService . switch_tenant ( account , tenant . id )
@ -679,6 +873,11 @@ class RegisterService:
redis_client . setex ( cls . _get_invitation_token_key ( token ) , expiry_hours * 60 * 60 , json . dumps ( invitation_data ) )
return token
@classmethod
def is_valid_invite_token ( cls , token : str ) - > bool :
data = redis_client . get ( cls . _get_invitation_token_key ( token ) )
return data is not None
@classmethod
def revoke_token ( cls , workspace_id : str , email : str , token : str ) :
if workspace_id and email :
@ -727,7 +926,9 @@ class RegisterService:
}
@classmethod
def _get_invitation_by_token ( cls , token : str , workspace_id : str , email : str ) - > Optional [ dict [ str , str ] ] :
def _get_invitation_by_token (
cls , token : str , workspace_id : Optional [ str ] = None , email : Optional [ str ] = None
) - > Optional [ dict [ str , 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 } "
Joe
1 year ago
committed by
GitHub