estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Feat/new login (#8120)

Browse Source 
Co-authored-by: douxc <douxc512@gmail.com>
Co-authored-by: Garfield Dai <dai.hai@foxmail.com>
This commit is contained in:
Joe
2024-10-21 10:03:40 +08:00
committed by GitHub
parent 2c0eaaec3d
commit 4fd2743efa
24 changed files with 1027 additions and 292 deletions
Download Patch File Download Diff File Expand all files Collapse all files

255
api/services/account_service.py
View File

@@ -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 AccountLoginError("Invalid email or password.")
raise AccountNotFoundError()
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.")
db.session.commit()
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 = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
db.session.add(ta)
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}"

4
api/services/errors/account.py
View File

@@ -13,6 +13,10 @@ class AccountLoginError(BaseServiceError):
pass
class AccountPasswordError(BaseServiceError):
pass
class AccountNotLinkTenantError(BaseServiceError):
pass

9
api/services/errors/workspace.py Normal file
View File

@@ -0,0 +1,9 @@
from services.errors.base import BaseServiceError
class WorkSpaceNotAllowedCreateError(BaseServiceError):
pass
class WorkSpaceNotFoundError(BaseServiceError):
pass

36
api/services/feature_service.py
View File

@@ -42,6 +42,11 @@ class SystemFeatureModel(BaseModel):
sso_enforced_for_web: bool = False
sso_enforced_for_web_protocol: str = ""
enable_web_sso_switch_component: bool = False
enable_email_code_login: bool = False
enable_email_password_login: bool = True
enable_social_oauth_login: bool = False
is_allow_register: bool = False
is_allow_create_workspace: bool = False
class FeatureService:
@@ -60,12 +65,23 @@ class FeatureService:
def get_system_features(cls) -> SystemFeatureModel:
system_features = SystemFeatureModel()
cls._fulfill_system_params_from_env(system_features)
if dify_config.ENTERPRISE_ENABLED:
system_features.enable_web_sso_switch_component = True
cls._fulfill_params_from_enterprise(system_features)
return system_features
@classmethod
def _fulfill_system_params_from_env(cls, system_features: SystemFeatureModel):
system_features.enable_email_code_login = dify_config.ENABLE_EMAIL_CODE_LOGIN
system_features.enable_email_password_login = dify_config.ENABLE_EMAIL_PASSWORD_LOGIN
system_features.enable_social_oauth_login = dify_config.ENABLE_SOCIAL_OAUTH_LOGIN
system_features.is_allow_register = dify_config.ALLOW_REGISTER
system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE
@classmethod
def _fulfill_params_from_env(cls, features: FeatureModel):
features.can_replace_logo = dify_config.CAN_REPLACE_LOGO
@@ -113,7 +129,19 @@ class FeatureService:
def _fulfill_params_from_enterprise(cls, features):
enterprise_info = EnterpriseService.get_info()
features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"]
features.sso_enforced_for_signin_protocol = enterprise_info["sso_enforced_for_signin_protocol"]
features.sso_enforced_for_web = enterprise_info["sso_enforced_for_web"]
features.sso_enforced_for_web_protocol = enterprise_info["sso_enforced_for_web_protocol"]
if "sso_enforced_for_signin" in enterprise_info:
features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"]
if "sso_enforced_for_signin_protocol" in enterprise_info:
features.sso_enforced_for_signin_protocol = enterprise_info["sso_enforced_for_signin_protocol"]
if "sso_enforced_for_web" in enterprise_info:
features.sso_enforced_for_web = enterprise_info["sso_enforced_for_web"]
if "sso_enforced_for_web_protocol" in enterprise_info:
features.sso_enforced_for_web_protocol = enterprise_info["sso_enforced_for_web_protocol"]
if "enable_email_code_login" in enterprise_info:
features.enable_email_code_login = enterprise_info["enable_email_code_login"]
if "enable_email_password_login" in enterprise_info:
features.enable_email_password_login = enterprise_info["enable_email_password_login"]
if "is_allow_register" in enterprise_info:
features.is_allow_register = enterprise_info["is_allow_register"]
if "is_allow_create_workspace" in enterprise_info:
features.is_allow_create_workspace = enterprise_info["is_allow_create_workspace"]