Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

account_service.py 48KB

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