Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

account_service.py 50KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384
  1. import base64
  2. import json
  3. import logging
  4. import secrets
  5. import uuid
  6. from datetime import UTC, datetime, timedelta
  7. from hashlib import sha256
  8. from typing import Any, Optional, cast
  9. from pydantic import BaseModel
  10. from sqlalchemy import func
  11. from sqlalchemy.orm import Session
  12. from werkzeug.exceptions import Unauthorized
  13. from configs import dify_config
  14. from constants.languages import language_timezone_mapping, languages
  15. from events.tenant_event import tenant_was_created
  16. from extensions.ext_database import db
  17. from extensions.ext_redis import redis_client, redis_fallback
  18. from libs.datetime_utils import naive_utc_now
  19. from libs.helper import RateLimiter, TokenManager
  20. from libs.passport import PassportService
  21. from libs.password import compare_password, hash_password, valid_password
  22. from libs.rsa import generate_key_pair
  23. from models.account import (
  24. Account,
  25. AccountIntegrate,
  26. AccountStatus,
  27. Tenant,
  28. TenantAccountJoin,
  29. TenantAccountRole,
  30. TenantPluginAutoUpgradeStrategy,
  31. TenantStatus,
  32. )
  33. from models.model import DifySetup
  34. from services.billing_service import BillingService
  35. from services.errors.account import (
  36. AccountAlreadyInTenantError,
  37. AccountLoginError,
  38. AccountNotFoundError,
  39. AccountNotLinkTenantError,
  40. AccountPasswordError,
  41. AccountRegisterError,
  42. CannotOperateSelfError,
  43. CurrentPasswordIncorrectError,
  44. InvalidActionError,
  45. LinkAccountIntegrateError,
  46. MemberNotInTenantError,
  47. NoPermissionError,
  48. RoleAlreadyAssignedError,
  49. TenantNotFoundError,
  50. )
  51. from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
  52. from services.feature_service import FeatureService
  53. from tasks.delete_account_task import delete_account_task
  54. from tasks.mail_account_deletion_task import send_account_deletion_verification_code
  55. from tasks.mail_change_mail_task import (
  56. send_change_mail_completed_notification_task,
  57. send_change_mail_task,
  58. )
  59. from tasks.mail_email_code_login import send_email_code_login_mail_task
  60. from tasks.mail_invite_member_task import send_invite_member_mail_task
  61. from tasks.mail_owner_transfer_task import (
  62. send_new_owner_transfer_notify_email_task,
  63. send_old_owner_transfer_notify_email_task,
  64. send_owner_transfer_confirm_task,
  65. )
  66. from tasks.mail_reset_password_task import send_reset_password_mail_task
  67. logger = logging.getLogger(__name__)
  68. class TokenPair(BaseModel):
  69. access_token: str
  70. refresh_token: str
  71. REFRESH_TOKEN_PREFIX = "refresh_token:"
  72. ACCOUNT_REFRESH_TOKEN_PREFIX = "account_refresh_token:"
  73. REFRESH_TOKEN_EXPIRY = timedelta(days=dify_config.REFRESH_TOKEN_EXPIRE_DAYS)
  74. class AccountService:
  75. reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=1, time_window=60 * 1)
  76. email_code_login_rate_limiter = RateLimiter(
  77. prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1
  78. )
  79. email_code_account_deletion_rate_limiter = RateLimiter(
  80. prefix="email_code_account_deletion_rate_limit", max_attempts=1, time_window=60 * 1
  81. )
  82. change_email_rate_limiter = RateLimiter(prefix="change_email_rate_limit", max_attempts=1, time_window=60 * 1)
  83. owner_transfer_rate_limiter = RateLimiter(prefix="owner_transfer_rate_limit", max_attempts=1, time_window=60 * 1)
  84. LOGIN_MAX_ERROR_LIMITS = 5
  85. FORGOT_PASSWORD_MAX_ERROR_LIMITS = 5
  86. CHANGE_EMAIL_MAX_ERROR_LIMITS = 5
  87. OWNER_TRANSFER_MAX_ERROR_LIMITS = 5
  88. @staticmethod
  89. def _get_refresh_token_key(refresh_token: str) -> str:
  90. return f"{REFRESH_TOKEN_PREFIX}{refresh_token}"
  91. @staticmethod
  92. def _get_account_refresh_token_key(account_id: str) -> str:
  93. return f"{ACCOUNT_REFRESH_TOKEN_PREFIX}{account_id}"
  94. @staticmethod
  95. def _store_refresh_token(refresh_token: str, account_id: str) -> None:
  96. redis_client.setex(AccountService._get_refresh_token_key(refresh_token), REFRESH_TOKEN_EXPIRY, account_id)
  97. redis_client.setex(
  98. AccountService._get_account_refresh_token_key(account_id), REFRESH_TOKEN_EXPIRY, refresh_token
  99. )
  100. @staticmethod
  101. def _delete_refresh_token(refresh_token: str, account_id: str) -> None:
  102. redis_client.delete(AccountService._get_refresh_token_key(refresh_token))
  103. redis_client.delete(AccountService._get_account_refresh_token_key(account_id))
  104. @staticmethod
  105. def load_user(user_id: str) -> None | Account:
  106. account = db.session.query(Account).filter_by(id=user_id).first()
  107. if not account:
  108. return None
  109. if account.status == AccountStatus.BANNED.value:
  110. raise Unauthorized("Account is banned.")
  111. current_tenant = db.session.query(TenantAccountJoin).filter_by(account_id=account.id, current=True).first()
  112. if current_tenant:
  113. account.set_tenant_id(current_tenant.tenant_id)
  114. else:
  115. available_ta = (
  116. db.session.query(TenantAccountJoin)
  117. .filter_by(account_id=account.id)
  118. .order_by(TenantAccountJoin.id.asc())
  119. .first()
  120. )
  121. if not available_ta:
  122. return None
  123. account.set_tenant_id(available_ta.tenant_id)
  124. available_ta.current = True
  125. db.session.commit()
  126. if naive_utc_now() - account.last_active_at > timedelta(minutes=10):
  127. account.last_active_at = naive_utc_now()
  128. db.session.commit()
  129. # NOTE: make sure account is accessible outside of a db session
  130. # This ensures that it will work correctly after upgrading to Flask version 3.1.2
  131. db.session.refresh(account)
  132. db.session.close()
  133. return account
  134. @staticmethod
  135. def get_account_jwt_token(account: Account) -> str:
  136. exp_dt = datetime.now(UTC) + timedelta(minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES)
  137. exp = int(exp_dt.timestamp())
  138. payload = {
  139. "user_id": account.id,
  140. "exp": exp,
  141. "iss": dify_config.EDITION,
  142. "sub": "Console API Passport",
  143. }
  144. token: str = PassportService().issue(payload)
  145. return token
  146. @staticmethod
  147. def authenticate(email: str, password: str, invite_token: Optional[str] = None) -> Account:
  148. """authenticate account with email and password"""
  149. account = db.session.query(Account).filter_by(email=email).first()
  150. if not account:
  151. raise AccountNotFoundError()
  152. if account.status == AccountStatus.BANNED.value:
  153. raise AccountLoginError("Account is banned.")
  154. if password and invite_token and account.password is None:
  155. # if invite_token is valid, set password and password_salt
  156. salt = secrets.token_bytes(16)
  157. base64_salt = base64.b64encode(salt).decode()
  158. password_hashed = hash_password(password, salt)
  159. base64_password_hashed = base64.b64encode(password_hashed).decode()
  160. account.password = base64_password_hashed
  161. account.password_salt = base64_salt
  162. if account.password is None or not compare_password(password, account.password, account.password_salt):
  163. raise AccountPasswordError("Invalid email or password.")
  164. if account.status == AccountStatus.PENDING.value:
  165. account.status = AccountStatus.ACTIVE.value
  166. account.initialized_at = naive_utc_now()
  167. db.session.commit()
  168. return account
  169. @staticmethod
  170. def update_account_password(account, password, new_password):
  171. """update account password"""
  172. if account.password and not compare_password(password, account.password, account.password_salt):
  173. raise CurrentPasswordIncorrectError("Current password is incorrect.")
  174. # may be raised
  175. valid_password(new_password)
  176. # generate password salt
  177. salt = secrets.token_bytes(16)
  178. base64_salt = base64.b64encode(salt).decode()
  179. # encrypt password with salt
  180. password_hashed = hash_password(new_password, salt)
  181. base64_password_hashed = base64.b64encode(password_hashed).decode()
  182. account.password = base64_password_hashed
  183. account.password_salt = base64_salt
  184. db.session.commit()
  185. return account
  186. @staticmethod
  187. def create_account(
  188. email: str,
  189. name: str,
  190. interface_language: str,
  191. password: Optional[str] = None,
  192. interface_theme: str = "light",
  193. is_setup: Optional[bool] = False,
  194. ) -> Account:
  195. """create account"""
  196. if not FeatureService.get_system_features().is_allow_register and not is_setup:
  197. from controllers.console.error import AccountNotFound
  198. raise AccountNotFound()
  199. if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
  200. raise AccountRegisterError(
  201. description=(
  202. "This email account has been deleted within the past "
  203. "30 days and is temporarily unavailable for new account registration"
  204. )
  205. )
  206. account = Account()
  207. account.email = email
  208. account.name = name
  209. if password:
  210. # generate password salt
  211. salt = secrets.token_bytes(16)
  212. base64_salt = base64.b64encode(salt).decode()
  213. # encrypt password with salt
  214. password_hashed = hash_password(password, salt)
  215. base64_password_hashed = base64.b64encode(password_hashed).decode()
  216. account.password = base64_password_hashed
  217. account.password_salt = base64_salt
  218. account.interface_language = interface_language
  219. account.interface_theme = interface_theme
  220. # Set timezone based on language
  221. account.timezone = language_timezone_mapping.get(interface_language, "UTC")
  222. db.session.add(account)
  223. db.session.commit()
  224. return account
  225. @staticmethod
  226. def create_account_and_tenant(
  227. email: str, name: str, interface_language: str, password: Optional[str] = None
  228. ) -> Account:
  229. """create account"""
  230. account = AccountService.create_account(
  231. email=email, name=name, interface_language=interface_language, password=password
  232. )
  233. TenantService.create_owner_tenant_if_not_exist(account=account)
  234. return account
  235. @staticmethod
  236. def generate_account_deletion_verification_code(account: Account) -> tuple[str, str]:
  237. code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)])
  238. token = TokenManager.generate_token(
  239. account=account, token_type="account_deletion", additional_data={"code": code}
  240. )
  241. return token, code
  242. @classmethod
  243. def send_account_deletion_verification_email(cls, account: Account, code: str):
  244. email = account.email
  245. if cls.email_code_account_deletion_rate_limiter.is_rate_limited(email):
  246. from controllers.console.auth.error import EmailCodeAccountDeletionRateLimitExceededError
  247. raise EmailCodeAccountDeletionRateLimitExceededError()
  248. send_account_deletion_verification_code.delay(to=email, code=code)
  249. cls.email_code_account_deletion_rate_limiter.increment_rate_limit(email)
  250. @staticmethod
  251. def verify_account_deletion_code(token: str, code: str) -> bool:
  252. token_data = TokenManager.get_token_data(token, "account_deletion")
  253. if token_data is None:
  254. return False
  255. if token_data["code"] != code:
  256. return False
  257. return True
  258. @staticmethod
  259. def delete_account(account: Account) -> None:
  260. """Delete account. This method only adds a task to the queue for deletion."""
  261. delete_account_task.delay(account.id)
  262. @staticmethod
  263. def link_account_integrate(provider: str, open_id: str, account: Account) -> None:
  264. """Link account integrate"""
  265. try:
  266. # Query whether there is an existing binding record for the same provider
  267. account_integrate: Optional[AccountIntegrate] = (
  268. db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider=provider).first()
  269. )
  270. if account_integrate:
  271. # If it exists, update the record
  272. account_integrate.open_id = open_id
  273. account_integrate.encrypted_token = "" # todo
  274. account_integrate.updated_at = naive_utc_now()
  275. else:
  276. # If it does not exist, create a new record
  277. account_integrate = AccountIntegrate(
  278. account_id=account.id, provider=provider, open_id=open_id, encrypted_token=""
  279. )
  280. db.session.add(account_integrate)
  281. db.session.commit()
  282. logger.info("Account %s linked %s account %s.", account.id, provider, open_id)
  283. except Exception as e:
  284. logger.exception("Failed to link %s account %s to Account %s", provider, open_id, account.id)
  285. raise LinkAccountIntegrateError("Failed to link account.") from e
  286. @staticmethod
  287. def close_account(account: Account) -> None:
  288. """Close account"""
  289. account.status = AccountStatus.CLOSED.value
  290. db.session.commit()
  291. @staticmethod
  292. def update_account(account, **kwargs):
  293. """Update account fields"""
  294. account = db.session.merge(account)
  295. for field, value in kwargs.items():
  296. if hasattr(account, field):
  297. setattr(account, field, value)
  298. else:
  299. raise AttributeError(f"Invalid field: {field}")
  300. db.session.commit()
  301. return account
  302. @staticmethod
  303. def update_account_email(account: Account, email: str) -> Account:
  304. """Update account email"""
  305. account.email = email
  306. account_integrate = db.session.query(AccountIntegrate).filter_by(account_id=account.id).first()
  307. if account_integrate:
  308. db.session.delete(account_integrate)
  309. db.session.add(account)
  310. db.session.commit()
  311. return account
  312. @staticmethod
  313. def update_login_info(account: Account, *, ip_address: str) -> None:
  314. """Update last login time and ip"""
  315. account.last_login_at = naive_utc_now()
  316. account.last_login_ip = ip_address
  317. db.session.add(account)
  318. db.session.commit()
  319. @staticmethod
  320. def login(account: Account, *, ip_address: Optional[str] = None) -> TokenPair:
  321. if ip_address:
  322. AccountService.update_login_info(account=account, ip_address=ip_address)
  323. if account.status == AccountStatus.PENDING.value:
  324. account.status = AccountStatus.ACTIVE.value
  325. db.session.commit()
  326. access_token = AccountService.get_account_jwt_token(account=account)
  327. refresh_token = _generate_refresh_token()
  328. AccountService._store_refresh_token(refresh_token, account.id)
  329. return TokenPair(access_token=access_token, refresh_token=refresh_token)
  330. @staticmethod
  331. def logout(*, account: Account) -> None:
  332. refresh_token = redis_client.get(AccountService._get_account_refresh_token_key(account.id))
  333. if refresh_token:
  334. AccountService._delete_refresh_token(refresh_token.decode("utf-8"), account.id)
  335. @staticmethod
  336. def refresh_token(refresh_token: str) -> TokenPair:
  337. # Verify the refresh token
  338. account_id = redis_client.get(AccountService._get_refresh_token_key(refresh_token))
  339. if not account_id:
  340. raise ValueError("Invalid refresh token")
  341. account = AccountService.load_user(account_id.decode("utf-8"))
  342. if not account:
  343. raise ValueError("Invalid account")
  344. # Generate new access token and refresh token
  345. new_access_token = AccountService.get_account_jwt_token(account)
  346. new_refresh_token = _generate_refresh_token()
  347. AccountService._delete_refresh_token(refresh_token, account.id)
  348. AccountService._store_refresh_token(new_refresh_token, account.id)
  349. return TokenPair(access_token=new_access_token, refresh_token=new_refresh_token)
  350. @staticmethod
  351. def load_logged_in_account(*, account_id: str):
  352. return AccountService.load_user(account_id)
  353. @classmethod
  354. def send_reset_password_email(
  355. cls,
  356. account: Optional[Account] = None,
  357. email: Optional[str] = None,
  358. language: str = "en-US",
  359. ):
  360. account_email = account.email if account else email
  361. if account_email is None:
  362. raise ValueError("Email must be provided.")
  363. if cls.reset_password_rate_limiter.is_rate_limited(account_email):
  364. from controllers.console.auth.error import PasswordResetRateLimitExceededError
  365. raise PasswordResetRateLimitExceededError()
  366. code, token = cls.generate_reset_password_token(account_email, account)
  367. send_reset_password_mail_task.delay(
  368. language=language,
  369. to=account_email,
  370. code=code,
  371. )
  372. cls.reset_password_rate_limiter.increment_rate_limit(account_email)
  373. return token
  374. @classmethod
  375. def send_change_email_email(
  376. cls,
  377. account: Optional[Account] = None,
  378. email: Optional[str] = None,
  379. old_email: Optional[str] = None,
  380. language: str = "en-US",
  381. phase: Optional[str] = None,
  382. ):
  383. account_email = account.email if account else email
  384. if account_email is None:
  385. raise ValueError("Email must be provided.")
  386. if not phase:
  387. raise ValueError("phase must be provided.")
  388. if cls.change_email_rate_limiter.is_rate_limited(account_email):
  389. from controllers.console.auth.error import EmailChangeRateLimitExceededError
  390. raise EmailChangeRateLimitExceededError()
  391. code, token = cls.generate_change_email_token(account_email, account, old_email=old_email)
  392. send_change_mail_task.delay(
  393. language=language,
  394. to=account_email,
  395. code=code,
  396. phase=phase,
  397. )
  398. cls.change_email_rate_limiter.increment_rate_limit(account_email)
  399. return token
  400. @classmethod
  401. def send_change_email_completed_notify_email(
  402. cls,
  403. account: Optional[Account] = None,
  404. email: Optional[str] = None,
  405. language: str = "en-US",
  406. ):
  407. account_email = account.email if account else email
  408. if account_email is None:
  409. raise ValueError("Email must be provided.")
  410. send_change_mail_completed_notification_task.delay(
  411. language=language,
  412. to=account_email,
  413. )
  414. @classmethod
  415. def send_owner_transfer_email(
  416. cls,
  417. account: Optional[Account] = None,
  418. email: Optional[str] = None,
  419. language: str = "en-US",
  420. workspace_name: Optional[str] = "",
  421. ):
  422. account_email = account.email if account else email
  423. if account_email is None:
  424. raise ValueError("Email must be provided.")
  425. if cls.owner_transfer_rate_limiter.is_rate_limited(account_email):
  426. from controllers.console.auth.error import OwnerTransferRateLimitExceededError
  427. raise OwnerTransferRateLimitExceededError()
  428. code, token = cls.generate_owner_transfer_token(account_email, account)
  429. workspace_name = workspace_name or ""
  430. send_owner_transfer_confirm_task.delay(
  431. language=language,
  432. to=account_email,
  433. code=code,
  434. workspace=workspace_name,
  435. )
  436. cls.owner_transfer_rate_limiter.increment_rate_limit(account_email)
  437. return token
  438. @classmethod
  439. def send_old_owner_transfer_notify_email(
  440. cls,
  441. account: Optional[Account] = None,
  442. email: Optional[str] = None,
  443. language: str = "en-US",
  444. workspace_name: Optional[str] = "",
  445. new_owner_email: str = "",
  446. ):
  447. account_email = account.email if account else email
  448. if account_email is None:
  449. raise ValueError("Email must be provided.")
  450. workspace_name = workspace_name or ""
  451. send_old_owner_transfer_notify_email_task.delay(
  452. language=language,
  453. to=account_email,
  454. workspace=workspace_name,
  455. new_owner_email=new_owner_email,
  456. )
  457. @classmethod
  458. def send_new_owner_transfer_notify_email(
  459. cls,
  460. account: Optional[Account] = None,
  461. email: Optional[str] = None,
  462. language: str = "en-US",
  463. workspace_name: Optional[str] = "",
  464. ):
  465. account_email = account.email if account else email
  466. if account_email is None:
  467. raise ValueError("Email must be provided.")
  468. workspace_name = workspace_name or ""
  469. send_new_owner_transfer_notify_email_task.delay(
  470. language=language,
  471. to=account_email,
  472. workspace=workspace_name,
  473. )
  474. @classmethod
  475. def generate_reset_password_token(
  476. cls,
  477. email: str,
  478. account: Optional[Account] = None,
  479. code: Optional[str] = None,
  480. additional_data: dict[str, Any] = {},
  481. ):
  482. if not code:
  483. code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)])
  484. additional_data["code"] = code
  485. token = TokenManager.generate_token(
  486. account=account, email=email, token_type="reset_password", additional_data=additional_data
  487. )
  488. return code, token
  489. @classmethod
  490. def generate_change_email_token(
  491. cls,
  492. email: str,
  493. account: Optional[Account] = None,
  494. code: Optional[str] = None,
  495. old_email: Optional[str] = None,
  496. additional_data: dict[str, Any] = {},
  497. ):
  498. if not code:
  499. code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)])
  500. additional_data["code"] = code
  501. additional_data["old_email"] = old_email
  502. token = TokenManager.generate_token(
  503. account=account, email=email, token_type="change_email", additional_data=additional_data
  504. )
  505. return code, token
  506. @classmethod
  507. def generate_owner_transfer_token(
  508. cls,
  509. email: str,
  510. account: Optional[Account] = None,
  511. code: Optional[str] = None,
  512. additional_data: dict[str, Any] = {},
  513. ):
  514. if not code:
  515. code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)])
  516. additional_data["code"] = code
  517. token = TokenManager.generate_token(
  518. account=account, email=email, token_type="owner_transfer", additional_data=additional_data
  519. )
  520. return code, token
  521. @classmethod
  522. def revoke_reset_password_token(cls, token: str):
  523. TokenManager.revoke_token(token, "reset_password")
  524. @classmethod
  525. def revoke_change_email_token(cls, token: str):
  526. TokenManager.revoke_token(token, "change_email")
  527. @classmethod
  528. def revoke_owner_transfer_token(cls, token: str):
  529. TokenManager.revoke_token(token, "owner_transfer")
  530. @classmethod
  531. def get_reset_password_data(cls, token: str) -> Optional[dict[str, Any]]:
  532. return TokenManager.get_token_data(token, "reset_password")
  533. @classmethod
  534. def get_change_email_data(cls, token: str) -> Optional[dict[str, Any]]:
  535. return TokenManager.get_token_data(token, "change_email")
  536. @classmethod
  537. def get_owner_transfer_data(cls, token: str) -> Optional[dict[str, Any]]:
  538. return TokenManager.get_token_data(token, "owner_transfer")
  539. @classmethod
  540. def send_email_code_login_email(
  541. cls,
  542. account: Optional[Account] = None,
  543. email: Optional[str] = None,
  544. language: str = "en-US",
  545. ):
  546. email = account.email if account else email
  547. if email is None:
  548. raise ValueError("Email must be provided.")
  549. if cls.email_code_login_rate_limiter.is_rate_limited(email):
  550. from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError
  551. raise EmailCodeLoginRateLimitExceededError()
  552. code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)])
  553. token = TokenManager.generate_token(
  554. account=account, email=email, token_type="email_code_login", additional_data={"code": code}
  555. )
  556. send_email_code_login_mail_task.delay(
  557. language=language,
  558. to=account.email if account else email,
  559. code=code,
  560. )
  561. cls.email_code_login_rate_limiter.increment_rate_limit(email)
  562. return token
  563. @classmethod
  564. def get_email_code_login_data(cls, token: str) -> Optional[dict[str, Any]]:
  565. return TokenManager.get_token_data(token, "email_code_login")
  566. @classmethod
  567. def revoke_email_code_login_token(cls, token: str):
  568. TokenManager.revoke_token(token, "email_code_login")
  569. @classmethod
  570. def get_user_through_email(cls, email: str):
  571. if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
  572. raise AccountRegisterError(
  573. description=(
  574. "This email account has been deleted within the past "
  575. "30 days and is temporarily unavailable for new account registration"
  576. )
  577. )
  578. account = db.session.query(Account).where(Account.email == email).first()
  579. if not account:
  580. return None
  581. if account.status == AccountStatus.BANNED.value:
  582. raise Unauthorized("Account is banned.")
  583. return account
  584. @classmethod
  585. def is_account_in_freeze(cls, email: str) -> bool:
  586. if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
  587. return True
  588. return False
  589. @staticmethod
  590. @redis_fallback(default_return=None)
  591. def add_login_error_rate_limit(email: str) -> None:
  592. key = f"login_error_rate_limit:{email}"
  593. count = redis_client.get(key)
  594. if count is None:
  595. count = 0
  596. count = int(count) + 1
  597. redis_client.setex(key, dify_config.LOGIN_LOCKOUT_DURATION, count)
  598. @staticmethod
  599. @redis_fallback(default_return=False)
  600. def is_login_error_rate_limit(email: str) -> bool:
  601. key = f"login_error_rate_limit:{email}"
  602. count = redis_client.get(key)
  603. if count is None:
  604. return False
  605. count = int(count)
  606. if count > AccountService.LOGIN_MAX_ERROR_LIMITS:
  607. return True
  608. return False
  609. @staticmethod
  610. @redis_fallback(default_return=None)
  611. def reset_login_error_rate_limit(email: str):
  612. key = f"login_error_rate_limit:{email}"
  613. redis_client.delete(key)
  614. @staticmethod
  615. @redis_fallback(default_return=None)
  616. def add_forgot_password_error_rate_limit(email: str) -> None:
  617. key = f"forgot_password_error_rate_limit:{email}"
  618. count = redis_client.get(key)
  619. if count is None:
  620. count = 0
  621. count = int(count) + 1
  622. redis_client.setex(key, dify_config.FORGOT_PASSWORD_LOCKOUT_DURATION, count)
  623. @staticmethod
  624. @redis_fallback(default_return=False)
  625. def is_forgot_password_error_rate_limit(email: str) -> bool:
  626. key = f"forgot_password_error_rate_limit:{email}"
  627. count = redis_client.get(key)
  628. if count is None:
  629. return False
  630. count = int(count)
  631. if count > AccountService.FORGOT_PASSWORD_MAX_ERROR_LIMITS:
  632. return True
  633. return False
  634. @staticmethod
  635. @redis_fallback(default_return=None)
  636. def reset_forgot_password_error_rate_limit(email: str):
  637. key = f"forgot_password_error_rate_limit:{email}"
  638. redis_client.delete(key)
  639. @staticmethod
  640. @redis_fallback(default_return=None)
  641. def add_change_email_error_rate_limit(email: str) -> None:
  642. key = f"change_email_error_rate_limit:{email}"
  643. count = redis_client.get(key)
  644. if count is None:
  645. count = 0
  646. count = int(count) + 1
  647. redis_client.setex(key, dify_config.CHANGE_EMAIL_LOCKOUT_DURATION, count)
  648. @staticmethod
  649. @redis_fallback(default_return=False)
  650. def is_change_email_error_rate_limit(email: str) -> bool:
  651. key = f"change_email_error_rate_limit:{email}"
  652. count = redis_client.get(key)
  653. if count is None:
  654. return False
  655. count = int(count)
  656. if count > AccountService.CHANGE_EMAIL_MAX_ERROR_LIMITS:
  657. return True
  658. return False
  659. @staticmethod
  660. @redis_fallback(default_return=None)
  661. def reset_change_email_error_rate_limit(email: str):
  662. key = f"change_email_error_rate_limit:{email}"
  663. redis_client.delete(key)
  664. @staticmethod
  665. @redis_fallback(default_return=None)
  666. def add_owner_transfer_error_rate_limit(email: str) -> None:
  667. key = f"owner_transfer_error_rate_limit:{email}"
  668. count = redis_client.get(key)
  669. if count is None:
  670. count = 0
  671. count = int(count) + 1
  672. redis_client.setex(key, dify_config.OWNER_TRANSFER_LOCKOUT_DURATION, count)
  673. @staticmethod
  674. @redis_fallback(default_return=False)
  675. def is_owner_transfer_error_rate_limit(email: str) -> bool:
  676. key = f"owner_transfer_error_rate_limit:{email}"
  677. count = redis_client.get(key)
  678. if count is None:
  679. return False
  680. count = int(count)
  681. if count > AccountService.OWNER_TRANSFER_MAX_ERROR_LIMITS:
  682. return True
  683. return False
  684. @staticmethod
  685. @redis_fallback(default_return=None)
  686. def reset_owner_transfer_error_rate_limit(email: str):
  687. key = f"owner_transfer_error_rate_limit:{email}"
  688. redis_client.delete(key)
  689. @staticmethod
  690. @redis_fallback(default_return=False)
  691. def is_email_send_ip_limit(ip_address: str):
  692. minute_key = f"email_send_ip_limit_minute:{ip_address}"
  693. freeze_key = f"email_send_ip_limit_freeze:{ip_address}"
  694. hour_limit_key = f"email_send_ip_limit_hour:{ip_address}"
  695. # check ip is frozen
  696. if redis_client.get(freeze_key):
  697. return True
  698. # check current minute count
  699. current_minute_count = redis_client.get(minute_key)
  700. if current_minute_count is None:
  701. current_minute_count = 0
  702. current_minute_count = int(current_minute_count)
  703. # check current hour count
  704. if current_minute_count > dify_config.EMAIL_SEND_IP_LIMIT_PER_MINUTE:
  705. hour_limit_count = redis_client.get(hour_limit_key)
  706. if hour_limit_count is None:
  707. hour_limit_count = 0
  708. hour_limit_count = int(hour_limit_count)
  709. if hour_limit_count >= 1:
  710. redis_client.setex(freeze_key, 60 * 60, 1)
  711. return True
  712. else:
  713. redis_client.setex(hour_limit_key, 60 * 10, hour_limit_count + 1) # first time limit 10 minutes
  714. # add hour limit count
  715. redis_client.incr(hour_limit_key)
  716. redis_client.expire(hour_limit_key, 60 * 60)
  717. return True
  718. redis_client.setex(minute_key, 60, current_minute_count + 1)
  719. redis_client.expire(minute_key, 60)
  720. return False
  721. @staticmethod
  722. def check_email_unique(email: str) -> bool:
  723. return db.session.query(Account).filter_by(email=email).first() is None
  724. class TenantService:
  725. @staticmethod
  726. def create_tenant(name: str, is_setup: Optional[bool] = False, is_from_dashboard: Optional[bool] = False) -> Tenant:
  727. """Create tenant"""
  728. if (
  729. not FeatureService.get_system_features().is_allow_create_workspace
  730. and not is_setup
  731. and not is_from_dashboard
  732. ):
  733. from controllers.console.error import NotAllowedCreateWorkspace
  734. raise NotAllowedCreateWorkspace()
  735. tenant = Tenant(name=name)
  736. db.session.add(tenant)
  737. db.session.commit()
  738. plugin_upgrade_strategy = TenantPluginAutoUpgradeStrategy(
  739. tenant_id=tenant.id,
  740. strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY,
  741. upgrade_time_of_day=0,
  742. upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
  743. exclude_plugins=[],
  744. include_plugins=[],
  745. )
  746. db.session.add(plugin_upgrade_strategy)
  747. db.session.commit()
  748. tenant.encrypt_public_key = generate_key_pair(tenant.id)
  749. db.session.commit()
  750. return tenant
  751. @staticmethod
  752. def create_owner_tenant_if_not_exist(
  753. account: Account, name: Optional[str] = None, is_setup: Optional[bool] = False
  754. ):
  755. """Check if user have a workspace or not"""
  756. available_ta = (
  757. db.session.query(TenantAccountJoin)
  758. .filter_by(account_id=account.id)
  759. .order_by(TenantAccountJoin.id.asc())
  760. .first()
  761. )
  762. if available_ta:
  763. return
  764. """Create owner tenant if not exist"""
  765. if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup:
  766. raise WorkSpaceNotAllowedCreateError()
  767. workspaces = FeatureService.get_system_features().license.workspaces
  768. if not workspaces.is_available():
  769. raise WorkspacesLimitExceededError()
  770. if name:
  771. tenant = TenantService.create_tenant(name=name, is_setup=is_setup)
  772. else:
  773. tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup)
  774. TenantService.create_tenant_member(tenant, account, role="owner")
  775. account.current_tenant = tenant
  776. db.session.commit()
  777. tenant_was_created.send(tenant)
  778. @staticmethod
  779. def create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin:
  780. """Create tenant member"""
  781. if role == TenantAccountRole.OWNER.value:
  782. if TenantService.has_roles(tenant, [TenantAccountRole.OWNER]):
  783. logger.error("Tenant %s has already an owner.", tenant.id)
  784. raise Exception("Tenant already has an owner.")
  785. ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()
  786. if ta:
  787. ta.role = role
  788. else:
  789. ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
  790. db.session.add(ta)
  791. db.session.commit()
  792. return ta
  793. @staticmethod
  794. def get_join_tenants(account: Account) -> list[Tenant]:
  795. """Get account join tenants"""
  796. return (
  797. db.session.query(Tenant)
  798. .join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
  799. .where(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL)
  800. .all()
  801. )
  802. @staticmethod
  803. def get_current_tenant_by_account(account: Account):
  804. """Get tenant by account and add the role"""
  805. tenant = account.current_tenant
  806. if not tenant:
  807. raise TenantNotFoundError("Tenant not found.")
  808. ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()
  809. if ta:
  810. tenant.role = ta.role
  811. else:
  812. raise TenantNotFoundError("Tenant not found for the account.")
  813. return tenant
  814. @staticmethod
  815. def switch_tenant(account: Account, tenant_id: Optional[str] = None) -> None:
  816. """Switch the current workspace for the account"""
  817. # Ensure tenant_id is provided
  818. if tenant_id is None:
  819. raise ValueError("Tenant ID must be provided.")
  820. tenant_account_join = (
  821. db.session.query(TenantAccountJoin)
  822. .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
  823. .where(
  824. TenantAccountJoin.account_id == account.id,
  825. TenantAccountJoin.tenant_id == tenant_id,
  826. Tenant.status == TenantStatus.NORMAL,
  827. )
  828. .first()
  829. )
  830. if not tenant_account_join:
  831. raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.")
  832. else:
  833. db.session.query(TenantAccountJoin).where(
  834. TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id
  835. ).update({"current": False})
  836. tenant_account_join.current = True
  837. # Set the current tenant for the account
  838. account.set_tenant_id(tenant_account_join.tenant_id)
  839. db.session.commit()
  840. @staticmethod
  841. def get_tenant_members(tenant: Tenant) -> list[Account]:
  842. """Get tenant members"""
  843. query = (
  844. db.session.query(Account, TenantAccountJoin.role)
  845. .select_from(Account)
  846. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  847. .where(TenantAccountJoin.tenant_id == tenant.id)
  848. )
  849. # Initialize an empty list to store the updated accounts
  850. updated_accounts = []
  851. for account, role in query:
  852. account.role = role
  853. updated_accounts.append(account)
  854. return updated_accounts
  855. @staticmethod
  856. def get_dataset_operator_members(tenant: Tenant) -> list[Account]:
  857. """Get dataset admin members"""
  858. query = (
  859. db.session.query(Account, TenantAccountJoin.role)
  860. .select_from(Account)
  861. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  862. .where(TenantAccountJoin.tenant_id == tenant.id)
  863. .where(TenantAccountJoin.role == "dataset_operator")
  864. )
  865. # Initialize an empty list to store the updated accounts
  866. updated_accounts = []
  867. for account, role in query:
  868. account.role = role
  869. updated_accounts.append(account)
  870. return updated_accounts
  871. @staticmethod
  872. def has_roles(tenant: Tenant, roles: list[TenantAccountRole]) -> bool:
  873. """Check if user has any of the given roles for a tenant"""
  874. if not all(isinstance(role, TenantAccountRole) for role in roles):
  875. raise ValueError("all roles must be TenantAccountRole")
  876. return (
  877. db.session.query(TenantAccountJoin)
  878. .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles]))
  879. .first()
  880. is not None
  881. )
  882. @staticmethod
  883. def get_user_role(account: Account, tenant: Tenant) -> Optional[TenantAccountRole]:
  884. """Get the role of the current account for a given tenant"""
  885. join = (
  886. db.session.query(TenantAccountJoin)
  887. .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id)
  888. .first()
  889. )
  890. return TenantAccountRole(join.role) if join else None
  891. @staticmethod
  892. def get_tenant_count() -> int:
  893. """Get tenant count"""
  894. return cast(int, db.session.query(func.count(Tenant.id)).scalar())
  895. @staticmethod
  896. def check_member_permission(tenant: Tenant, operator: Account, member: Account | None, action: str) -> None:
  897. """Check member permission"""
  898. perms = {
  899. "add": [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
  900. "remove": [TenantAccountRole.OWNER],
  901. "update": [TenantAccountRole.OWNER],
  902. }
  903. if action not in {"add", "remove", "update"}:
  904. raise InvalidActionError("Invalid action.")
  905. if member:
  906. if operator.id == member.id:
  907. raise CannotOperateSelfError("Cannot operate self.")
  908. ta_operator = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=operator.id).first()
  909. if not ta_operator or ta_operator.role not in perms[action]:
  910. raise NoPermissionError(f"No permission to {action} member.")
  911. @staticmethod
  912. def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account) -> None:
  913. """Remove member from tenant"""
  914. if operator.id == account.id:
  915. raise CannotOperateSelfError("Cannot operate self.")
  916. TenantService.check_member_permission(tenant, operator, account, "remove")
  917. ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()
  918. if not ta:
  919. raise MemberNotInTenantError("Member not in tenant.")
  920. db.session.delete(ta)
  921. db.session.commit()
  922. @staticmethod
  923. def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account) -> None:
  924. """Update member role"""
  925. TenantService.check_member_permission(tenant, operator, member, "update")
  926. target_member_join = (
  927. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member.id).first()
  928. )
  929. if not target_member_join:
  930. raise MemberNotInTenantError("Member not in tenant.")
  931. if target_member_join.role == new_role:
  932. raise RoleAlreadyAssignedError("The provided role is already assigned to the member.")
  933. if new_role == "owner":
  934. # Find the current owner and change their role to 'admin'
  935. current_owner_join = (
  936. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, role="owner").first()
  937. )
  938. if current_owner_join:
  939. current_owner_join.role = "admin"
  940. # Update the role of the target member
  941. target_member_join.role = new_role
  942. db.session.commit()
  943. @staticmethod
  944. def get_custom_config(tenant_id: str) -> dict:
  945. tenant = db.get_or_404(Tenant, tenant_id)
  946. return tenant.custom_config_dict
  947. @staticmethod
  948. def is_owner(account: Account, tenant: Tenant) -> bool:
  949. return TenantService.get_user_role(account, tenant) == TenantAccountRole.OWNER
  950. @staticmethod
  951. def is_member(account: Account, tenant: Tenant) -> bool:
  952. """Check if the account is a member of the tenant"""
  953. return TenantService.get_user_role(account, tenant) is not None
  954. class RegisterService:
  955. @classmethod
  956. def _get_invitation_token_key(cls, token: str) -> str:
  957. return f"member_invite:token:{token}"
  958. @classmethod
  959. def setup(cls, email: str, name: str, password: str, ip_address: str) -> None:
  960. """
  961. Setup dify
  962. :param email: email
  963. :param name: username
  964. :param password: password
  965. :param ip_address: ip address
  966. """
  967. try:
  968. # Register
  969. account = AccountService.create_account(
  970. email=email,
  971. name=name,
  972. interface_language=languages[0],
  973. password=password,
  974. is_setup=True,
  975. )
  976. account.last_login_ip = ip_address
  977. account.initialized_at = naive_utc_now()
  978. TenantService.create_owner_tenant_if_not_exist(account=account, is_setup=True)
  979. dify_setup = DifySetup(version=dify_config.project.version)
  980. db.session.add(dify_setup)
  981. db.session.commit()
  982. except Exception as e:
  983. db.session.query(DifySetup).delete()
  984. db.session.query(TenantAccountJoin).delete()
  985. db.session.query(Account).delete()
  986. db.session.query(Tenant).delete()
  987. db.session.commit()
  988. logger.exception("Setup account failed, email: %s, name: %s", email, name)
  989. raise ValueError(f"Setup failed: {e}")
  990. @classmethod
  991. def register(
  992. cls,
  993. email,
  994. name,
  995. password: Optional[str] = None,
  996. open_id: Optional[str] = None,
  997. provider: Optional[str] = None,
  998. language: Optional[str] = None,
  999. status: Optional[AccountStatus] = None,
  1000. is_setup: Optional[bool] = False,
  1001. create_workspace_required: Optional[bool] = True,
  1002. ) -> Account:
  1003. db.session.begin_nested()
  1004. """Register account"""
  1005. try:
  1006. account = AccountService.create_account(
  1007. email=email,
  1008. name=name,
  1009. interface_language=language or languages[0],
  1010. password=password,
  1011. is_setup=is_setup,
  1012. )
  1013. account.status = AccountStatus.ACTIVE.value if not status else status.value
  1014. account.initialized_at = naive_utc_now()
  1015. if open_id is not None and provider is not None:
  1016. AccountService.link_account_integrate(provider, open_id, account)
  1017. if (
  1018. FeatureService.get_system_features().is_allow_create_workspace
  1019. and create_workspace_required
  1020. and FeatureService.get_system_features().license.workspaces.is_available()
  1021. ):
  1022. tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
  1023. TenantService.create_tenant_member(tenant, account, role="owner")
  1024. account.current_tenant = tenant
  1025. tenant_was_created.send(tenant)
  1026. db.session.commit()
  1027. except WorkSpaceNotAllowedCreateError:
  1028. db.session.rollback()
  1029. logger.exception("Register failed")
  1030. raise AccountRegisterError("Workspace is not allowed to create.")
  1031. except AccountRegisterError as are:
  1032. db.session.rollback()
  1033. logger.exception("Register failed")
  1034. raise are
  1035. except Exception as e:
  1036. db.session.rollback()
  1037. logger.exception("Register failed")
  1038. raise AccountRegisterError(f"Registration failed: {e}") from e
  1039. return account
  1040. @classmethod
  1041. def invite_new_member(
  1042. cls, tenant: Tenant, email: str, language: str, role: str = "normal", inviter: Account | None = None
  1043. ) -> str:
  1044. if not inviter:
  1045. raise ValueError("Inviter is required")
  1046. """Invite new member"""
  1047. with Session(db.engine) as session:
  1048. account = session.query(Account).filter_by(email=email).first()
  1049. if not account:
  1050. TenantService.check_member_permission(tenant, inviter, None, "add")
  1051. name = email.split("@")[0]
  1052. account = cls.register(
  1053. email=email, name=name, language=language, status=AccountStatus.PENDING, is_setup=True
  1054. )
  1055. # Create new tenant member for invited tenant
  1056. TenantService.create_tenant_member(tenant, account, role)
  1057. TenantService.switch_tenant(account, tenant.id)
  1058. else:
  1059. TenantService.check_member_permission(tenant, inviter, account, "add")
  1060. ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()
  1061. if not ta:
  1062. TenantService.create_tenant_member(tenant, account, role)
  1063. # Support resend invitation email when the account is pending status
  1064. if account.status != AccountStatus.PENDING.value:
  1065. raise AccountAlreadyInTenantError("Account already in tenant.")
  1066. token = cls.generate_invite_token(tenant, account)
  1067. language = account.interface_language or "en-US"
  1068. # send email
  1069. send_invite_member_mail_task.delay(
  1070. language=language,
  1071. to=email,
  1072. token=token,
  1073. inviter_name=inviter.name if inviter else "Dify",
  1074. workspace_name=tenant.name,
  1075. )
  1076. return token
  1077. @classmethod
  1078. def generate_invite_token(cls, tenant: Tenant, account: Account) -> str:
  1079. token = str(uuid.uuid4())
  1080. invitation_data = {
  1081. "account_id": account.id,
  1082. "email": account.email,
  1083. "workspace_id": tenant.id,
  1084. }
  1085. expiry_hours = dify_config.INVITE_EXPIRY_HOURS
  1086. redis_client.setex(cls._get_invitation_token_key(token), expiry_hours * 60 * 60, json.dumps(invitation_data))
  1087. return token
  1088. @classmethod
  1089. def is_valid_invite_token(cls, token: str) -> bool:
  1090. data = redis_client.get(cls._get_invitation_token_key(token))
  1091. return data is not None
  1092. @classmethod
  1093. def revoke_token(cls, workspace_id: str, email: str, token: str):
  1094. if workspace_id and email:
  1095. email_hash = sha256(email.encode()).hexdigest()
  1096. cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
  1097. redis_client.delete(cache_key)
  1098. else:
  1099. redis_client.delete(cls._get_invitation_token_key(token))
  1100. @classmethod
  1101. def get_invitation_if_token_valid(
  1102. cls, workspace_id: Optional[str], email: str, token: str
  1103. ) -> Optional[dict[str, Any]]:
  1104. invitation_data = cls._get_invitation_by_token(token, workspace_id, email)
  1105. if not invitation_data:
  1106. return None
  1107. tenant = (
  1108. db.session.query(Tenant)
  1109. .where(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal")
  1110. .first()
  1111. )
  1112. if not tenant:
  1113. return None
  1114. tenant_account = (
  1115. db.session.query(Account, TenantAccountJoin.role)
  1116. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  1117. .where(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id)
  1118. .first()
  1119. )
  1120. if not tenant_account:
  1121. return None
  1122. account = tenant_account[0]
  1123. if not account:
  1124. return None
  1125. if invitation_data["account_id"] != str(account.id):
  1126. return None
  1127. return {
  1128. "account": account,
  1129. "data": invitation_data,
  1130. "tenant": tenant,
  1131. }
  1132. @classmethod
  1133. def _get_invitation_by_token(
  1134. cls, token: str, workspace_id: Optional[str] = None, email: Optional[str] = None
  1135. ) -> Optional[dict[str, str]]:
  1136. if workspace_id is not None and email is not None:
  1137. email_hash = sha256(email.encode()).hexdigest()
  1138. cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
  1139. account_id = redis_client.get(cache_key)
  1140. if not account_id:
  1141. return None
  1142. return {
  1143. "account_id": account_id.decode("utf-8"),
  1144. "email": email,
  1145. "workspace_id": workspace_id,
  1146. }
  1147. else:
  1148. data = redis_client.get(cls._get_invitation_token_key(token))
  1149. if not data:
  1150. return None
  1151. invitation: dict = json.loads(data)
  1152. return invitation
  1153. def _generate_refresh_token(length: int = 64):
  1154. token = secrets.token_hex(length)
  1155. return token