Du kannst nicht mehr als 25 Themen auswählen Themen müssen mit entweder einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

account_service.py 48KB

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