account_service.py 39 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067
  1. import base64
  2. import json
  3. import logging
  4. import random
  5. import secrets
  6. import uuid
  7. from datetime import UTC, datetime, timedelta
  8. from hashlib import sha256
  9. from typing import Any, Optional, cast
  10. from pydantic import BaseModel
  11. from sqlalchemy import func
  12. from sqlalchemy.orm import Session
  13. from werkzeug.exceptions import Unauthorized
  14. from configs import dify_config
  15. from constants.languages import language_timezone_mapping, languages
  16. from events.tenant_event import tenant_was_created
  17. from extensions.ext_database import db
  18. from extensions.ext_redis import redis_client
  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
  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_email_code_login import send_email_code_login_mail_task
  55. from tasks.mail_invite_member_task import send_invite_member_mail_task
  56. from tasks.mail_reset_password_task import send_reset_password_mail_task
  57. class TokenPair(BaseModel):
  58. access_token: str
  59. refresh_token: str
  60. REFRESH_TOKEN_PREFIX = "refresh_token:"
  61. ACCOUNT_REFRESH_TOKEN_PREFIX = "account_refresh_token:"
  62. REFRESH_TOKEN_EXPIRY = timedelta(days=dify_config.REFRESH_TOKEN_EXPIRE_DAYS)
  63. class AccountService:
  64. reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=1, time_window=60 * 1)
  65. email_code_login_rate_limiter = RateLimiter(
  66. prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1
  67. )
  68. email_code_account_deletion_rate_limiter = RateLimiter(
  69. prefix="email_code_account_deletion_rate_limit", max_attempts=1, time_window=60 * 1
  70. )
  71. LOGIN_MAX_ERROR_LIMITS = 5
  72. FORGOT_PASSWORD_MAX_ERROR_LIMITS = 5
  73. @staticmethod
  74. def _get_refresh_token_key(refresh_token: str) -> str:
  75. return f"{REFRESH_TOKEN_PREFIX}{refresh_token}"
  76. @staticmethod
  77. def _get_account_refresh_token_key(account_id: str) -> str:
  78. return f"{ACCOUNT_REFRESH_TOKEN_PREFIX}{account_id}"
  79. @staticmethod
  80. def _store_refresh_token(refresh_token: str, account_id: str) -> None:
  81. redis_client.setex(AccountService._get_refresh_token_key(refresh_token), REFRESH_TOKEN_EXPIRY, account_id)
  82. redis_client.setex(
  83. AccountService._get_account_refresh_token_key(account_id), REFRESH_TOKEN_EXPIRY, refresh_token
  84. )
  85. @staticmethod
  86. def _delete_refresh_token(refresh_token: str, account_id: str) -> None:
  87. redis_client.delete(AccountService._get_refresh_token_key(refresh_token))
  88. redis_client.delete(AccountService._get_account_refresh_token_key(account_id))
  89. @staticmethod
  90. def load_user(user_id: str) -> None | Account:
  91. account = db.session.query(Account).filter_by(id=user_id).first()
  92. if not account:
  93. return None
  94. if account.status == AccountStatus.BANNED.value:
  95. raise Unauthorized("Account is banned.")
  96. current_tenant = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
  97. if current_tenant:
  98. account.current_tenant_id = current_tenant.tenant_id
  99. else:
  100. available_ta = (
  101. TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
  102. )
  103. if not available_ta:
  104. return None
  105. account.current_tenant_id = available_ta.tenant_id
  106. available_ta.current = True
  107. db.session.commit()
  108. if datetime.now(UTC).replace(tzinfo=None) - account.last_active_at > timedelta(minutes=10):
  109. account.last_active_at = datetime.now(UTC).replace(tzinfo=None)
  110. db.session.commit()
  111. return cast(Account, account)
  112. @staticmethod
  113. def get_account_jwt_token(account: Account) -> str:
  114. exp_dt = datetime.now(UTC) + timedelta(minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES)
  115. exp = int(exp_dt.timestamp())
  116. payload = {
  117. "user_id": account.id,
  118. "exp": exp,
  119. "iss": dify_config.EDITION,
  120. "sub": "Console API Passport",
  121. }
  122. token: str = PassportService().issue(payload)
  123. return token
  124. @staticmethod
  125. def authenticate(email: str, password: str, invite_token: Optional[str] = None) -> Account:
  126. """authenticate account with email and password"""
  127. account = db.session.query(Account).filter_by(email=email).first()
  128. if not account:
  129. raise AccountNotFoundError()
  130. if account.status == AccountStatus.BANNED.value:
  131. raise AccountLoginError("Account is banned.")
  132. if password and invite_token and account.password is None:
  133. # if invite_token is valid, set password and password_salt
  134. salt = secrets.token_bytes(16)
  135. base64_salt = base64.b64encode(salt).decode()
  136. password_hashed = hash_password(password, salt)
  137. base64_password_hashed = base64.b64encode(password_hashed).decode()
  138. account.password = base64_password_hashed
  139. account.password_salt = base64_salt
  140. if account.password is None or not compare_password(password, account.password, account.password_salt):
  141. raise AccountPasswordError("Invalid email or password.")
  142. if account.status == AccountStatus.PENDING.value:
  143. account.status = AccountStatus.ACTIVE.value
  144. account.initialized_at = datetime.now(UTC).replace(tzinfo=None)
  145. db.session.commit()
  146. return cast(Account, account)
  147. @staticmethod
  148. def update_account_password(account, password, new_password):
  149. """update account password"""
  150. if account.password and not compare_password(password, account.password, account.password_salt):
  151. raise CurrentPasswordIncorrectError("Current password is incorrect.")
  152. # may be raised
  153. valid_password(new_password)
  154. # generate password salt
  155. salt = secrets.token_bytes(16)
  156. base64_salt = base64.b64encode(salt).decode()
  157. # encrypt password with salt
  158. password_hashed = hash_password(new_password, salt)
  159. base64_password_hashed = base64.b64encode(password_hashed).decode()
  160. account.password = base64_password_hashed
  161. account.password_salt = base64_salt
  162. db.session.commit()
  163. return account
  164. @staticmethod
  165. def create_account(
  166. email: str,
  167. name: str,
  168. interface_language: str,
  169. password: Optional[str] = None,
  170. interface_theme: str = "light",
  171. is_setup: Optional[bool] = False,
  172. ) -> Account:
  173. """create account"""
  174. if not FeatureService.get_system_features().is_allow_register and not is_setup:
  175. from controllers.console.error import AccountNotFound
  176. raise AccountNotFound()
  177. if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
  178. raise AccountRegisterError(
  179. description=(
  180. "This email account has been deleted within the past "
  181. "30 days and is temporarily unavailable for new account registration"
  182. )
  183. )
  184. account = Account()
  185. account.email = email
  186. account.name = name
  187. if password:
  188. # generate password salt
  189. salt = secrets.token_bytes(16)
  190. base64_salt = base64.b64encode(salt).decode()
  191. # encrypt password with salt
  192. password_hashed = hash_password(password, salt)
  193. base64_password_hashed = base64.b64encode(password_hashed).decode()
  194. account.password = base64_password_hashed
  195. account.password_salt = base64_salt
  196. account.interface_language = interface_language
  197. account.interface_theme = interface_theme
  198. # Set timezone based on language
  199. account.timezone = language_timezone_mapping.get(interface_language, "UTC")
  200. db.session.add(account)
  201. db.session.commit()
  202. return account
  203. @staticmethod
  204. def create_account_and_tenant(
  205. email: str, name: str, interface_language: str, password: Optional[str] = None
  206. ) -> Account:
  207. """create account"""
  208. account = AccountService.create_account(
  209. email=email, name=name, interface_language=interface_language, password=password
  210. )
  211. TenantService.create_owner_tenant_if_not_exist(account=account)
  212. return account
  213. @staticmethod
  214. def generate_account_deletion_verification_code(account: Account) -> tuple[str, str]:
  215. code = "".join([str(random.randint(0, 9)) for _ in range(6)])
  216. token = TokenManager.generate_token(
  217. account=account, token_type="account_deletion", additional_data={"code": code}
  218. )
  219. return token, code
  220. @classmethod
  221. def send_account_deletion_verification_email(cls, account: Account, code: str):
  222. email = account.email
  223. if cls.email_code_account_deletion_rate_limiter.is_rate_limited(email):
  224. from controllers.console.auth.error import EmailCodeAccountDeletionRateLimitExceededError
  225. raise EmailCodeAccountDeletionRateLimitExceededError()
  226. send_account_deletion_verification_code.delay(to=email, code=code)
  227. cls.email_code_account_deletion_rate_limiter.increment_rate_limit(email)
  228. @staticmethod
  229. def verify_account_deletion_code(token: str, code: str) -> bool:
  230. token_data = TokenManager.get_token_data(token, "account_deletion")
  231. if token_data is None:
  232. return False
  233. if token_data["code"] != code:
  234. return False
  235. return True
  236. @staticmethod
  237. def delete_account(account: Account) -> None:
  238. """Delete account. This method only adds a task to the queue for deletion."""
  239. delete_account_task.delay(account.id)
  240. @staticmethod
  241. def link_account_integrate(provider: str, open_id: str, account: Account) -> None:
  242. """Link account integrate"""
  243. try:
  244. # Query whether there is an existing binding record for the same provider
  245. account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(
  246. account_id=account.id, provider=provider
  247. ).first()
  248. if account_integrate:
  249. # If it exists, update the record
  250. account_integrate.open_id = open_id
  251. account_integrate.encrypted_token = "" # todo
  252. account_integrate.updated_at = datetime.now(UTC).replace(tzinfo=None)
  253. else:
  254. # If it does not exist, create a new record
  255. account_integrate = AccountIntegrate(
  256. account_id=account.id, provider=provider, open_id=open_id, encrypted_token=""
  257. )
  258. db.session.add(account_integrate)
  259. db.session.commit()
  260. logging.info(f"Account {account.id} linked {provider} account {open_id}.")
  261. except Exception as e:
  262. logging.exception(f"Failed to link {provider} account {open_id} to Account {account.id}")
  263. raise LinkAccountIntegrateError("Failed to link account.") from e
  264. @staticmethod
  265. def close_account(account: Account) -> None:
  266. """Close account"""
  267. account.status = AccountStatus.CLOSED.value
  268. db.session.commit()
  269. @staticmethod
  270. def update_account(account, **kwargs):
  271. """Update account fields"""
  272. for field, value in kwargs.items():
  273. if hasattr(account, field):
  274. setattr(account, field, value)
  275. else:
  276. raise AttributeError(f"Invalid field: {field}")
  277. db.session.commit()
  278. return account
  279. @staticmethod
  280. def update_login_info(account: Account, *, ip_address: str) -> None:
  281. """Update last login time and ip"""
  282. account.last_login_at = datetime.now(UTC).replace(tzinfo=None)
  283. account.last_login_ip = ip_address
  284. db.session.add(account)
  285. db.session.commit()
  286. @staticmethod
  287. def login(account: Account, *, ip_address: Optional[str] = None) -> TokenPair:
  288. if ip_address:
  289. AccountService.update_login_info(account=account, ip_address=ip_address)
  290. if account.status == AccountStatus.PENDING.value:
  291. account.status = AccountStatus.ACTIVE.value
  292. db.session.commit()
  293. access_token = AccountService.get_account_jwt_token(account=account)
  294. refresh_token = _generate_refresh_token()
  295. AccountService._store_refresh_token(refresh_token, account.id)
  296. return TokenPair(access_token=access_token, refresh_token=refresh_token)
  297. @staticmethod
  298. def logout(*, account: Account) -> None:
  299. refresh_token = redis_client.get(AccountService._get_account_refresh_token_key(account.id))
  300. if refresh_token:
  301. AccountService._delete_refresh_token(refresh_token.decode("utf-8"), account.id)
  302. @staticmethod
  303. def refresh_token(refresh_token: str) -> TokenPair:
  304. # Verify the refresh token
  305. account_id = redis_client.get(AccountService._get_refresh_token_key(refresh_token))
  306. if not account_id:
  307. raise ValueError("Invalid refresh token")
  308. account = AccountService.load_user(account_id.decode("utf-8"))
  309. if not account:
  310. raise ValueError("Invalid account")
  311. # Generate new access token and refresh token
  312. new_access_token = AccountService.get_account_jwt_token(account)
  313. new_refresh_token = _generate_refresh_token()
  314. AccountService._delete_refresh_token(refresh_token, account.id)
  315. AccountService._store_refresh_token(new_refresh_token, account.id)
  316. return TokenPair(access_token=new_access_token, refresh_token=new_refresh_token)
  317. @staticmethod
  318. def load_logged_in_account(*, account_id: str):
  319. return AccountService.load_user(account_id)
  320. @classmethod
  321. def send_reset_password_email(
  322. cls,
  323. account: Optional[Account] = None,
  324. email: Optional[str] = None,
  325. language: Optional[str] = "en-US",
  326. ):
  327. account_email = account.email if account else email
  328. if account_email is None:
  329. raise ValueError("Email must be provided.")
  330. if cls.reset_password_rate_limiter.is_rate_limited(account_email):
  331. from controllers.console.auth.error import PasswordResetRateLimitExceededError
  332. raise PasswordResetRateLimitExceededError()
  333. code = "".join([str(random.randint(0, 9)) for _ in range(6)])
  334. token = TokenManager.generate_token(
  335. account=account, email=email, token_type="reset_password", additional_data={"code": code}
  336. )
  337. send_reset_password_mail_task.delay(
  338. language=language,
  339. to=account_email,
  340. code=code,
  341. )
  342. cls.reset_password_rate_limiter.increment_rate_limit(account_email)
  343. return token
  344. @classmethod
  345. def revoke_reset_password_token(cls, token: str):
  346. TokenManager.revoke_token(token, "reset_password")
  347. @classmethod
  348. def get_reset_password_data(cls, token: str) -> Optional[dict[str, Any]]:
  349. return TokenManager.get_token_data(token, "reset_password")
  350. @classmethod
  351. def send_email_code_login_email(
  352. cls, account: Optional[Account] = None, email: Optional[str] = None, language: Optional[str] = "en-US"
  353. ):
  354. email = account.email if account else email
  355. if email is None:
  356. raise ValueError("Email must be provided.")
  357. if cls.email_code_login_rate_limiter.is_rate_limited(email):
  358. from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError
  359. raise EmailCodeLoginRateLimitExceededError()
  360. code = "".join([str(random.randint(0, 9)) for _ in range(6)])
  361. token = TokenManager.generate_token(
  362. account=account, email=email, token_type="email_code_login", additional_data={"code": code}
  363. )
  364. send_email_code_login_mail_task.delay(
  365. language=language,
  366. to=account.email if account else email,
  367. code=code,
  368. )
  369. cls.email_code_login_rate_limiter.increment_rate_limit(email)
  370. return token
  371. @classmethod
  372. def get_email_code_login_data(cls, token: str) -> Optional[dict[str, Any]]:
  373. return TokenManager.get_token_data(token, "email_code_login")
  374. @classmethod
  375. def revoke_email_code_login_token(cls, token: str):
  376. TokenManager.revoke_token(token, "email_code_login")
  377. @classmethod
  378. def get_user_through_email(cls, email: str):
  379. if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):
  380. raise AccountRegisterError(
  381. description=(
  382. "This email account has been deleted within the past "
  383. "30 days and is temporarily unavailable for new account registration"
  384. )
  385. )
  386. account = db.session.query(Account).filter(Account.email == email).first()
  387. if not account:
  388. return None
  389. if account.status == AccountStatus.BANNED.value:
  390. raise Unauthorized("Account is banned.")
  391. return account
  392. @staticmethod
  393. def add_login_error_rate_limit(email: str) -> None:
  394. key = f"login_error_rate_limit:{email}"
  395. count = redis_client.get(key)
  396. if count is None:
  397. count = 0
  398. count = int(count) + 1
  399. redis_client.setex(key, dify_config.LOGIN_LOCKOUT_DURATION, count)
  400. @staticmethod
  401. def is_login_error_rate_limit(email: str) -> bool:
  402. key = f"login_error_rate_limit:{email}"
  403. count = redis_client.get(key)
  404. if count is None:
  405. return False
  406. count = int(count)
  407. if count > AccountService.LOGIN_MAX_ERROR_LIMITS:
  408. return True
  409. return False
  410. @staticmethod
  411. def reset_login_error_rate_limit(email: str):
  412. key = f"login_error_rate_limit:{email}"
  413. redis_client.delete(key)
  414. @staticmethod
  415. def add_forgot_password_error_rate_limit(email: str) -> None:
  416. key = f"forgot_password_error_rate_limit:{email}"
  417. count = redis_client.get(key)
  418. if count is None:
  419. count = 0
  420. count = int(count) + 1
  421. redis_client.setex(key, dify_config.FORGOT_PASSWORD_LOCKOUT_DURATION, count)
  422. @staticmethod
  423. def is_forgot_password_error_rate_limit(email: str) -> bool:
  424. key = f"forgot_password_error_rate_limit:{email}"
  425. count = redis_client.get(key)
  426. if count is None:
  427. return False
  428. count = int(count)
  429. if count > AccountService.FORGOT_PASSWORD_MAX_ERROR_LIMITS:
  430. return True
  431. return False
  432. @staticmethod
  433. def reset_forgot_password_error_rate_limit(email: str):
  434. key = f"forgot_password_error_rate_limit:{email}"
  435. redis_client.delete(key)
  436. @staticmethod
  437. def is_email_send_ip_limit(ip_address: str):
  438. minute_key = f"email_send_ip_limit_minute:{ip_address}"
  439. freeze_key = f"email_send_ip_limit_freeze:{ip_address}"
  440. hour_limit_key = f"email_send_ip_limit_hour:{ip_address}"
  441. # check ip is frozen
  442. if redis_client.get(freeze_key):
  443. return True
  444. # check current minute count
  445. current_minute_count = redis_client.get(minute_key)
  446. if current_minute_count is None:
  447. current_minute_count = 0
  448. current_minute_count = int(current_minute_count)
  449. # check current hour count
  450. if current_minute_count > dify_config.EMAIL_SEND_IP_LIMIT_PER_MINUTE:
  451. hour_limit_count = redis_client.get(hour_limit_key)
  452. if hour_limit_count is None:
  453. hour_limit_count = 0
  454. hour_limit_count = int(hour_limit_count)
  455. if hour_limit_count >= 1:
  456. redis_client.setex(freeze_key, 60 * 60, 1)
  457. return True
  458. else:
  459. redis_client.setex(hour_limit_key, 60 * 10, hour_limit_count + 1) # first time limit 10 minutes
  460. # add hour limit count
  461. redis_client.incr(hour_limit_key)
  462. redis_client.expire(hour_limit_key, 60 * 60)
  463. return True
  464. redis_client.setex(minute_key, 60, current_minute_count + 1)
  465. redis_client.expire(minute_key, 60)
  466. return False
  467. def _get_login_cache_key(*, account_id: str, token: str):
  468. return f"account_login:{account_id}:{token}"
  469. class TenantService:
  470. @staticmethod
  471. def create_tenant(name: str, is_setup: Optional[bool] = False, is_from_dashboard: Optional[bool] = False) -> Tenant:
  472. """Create tenant"""
  473. if (
  474. not FeatureService.get_system_features().is_allow_create_workspace
  475. and not is_setup
  476. and not is_from_dashboard
  477. ):
  478. from controllers.console.error import NotAllowedCreateWorkspace
  479. raise NotAllowedCreateWorkspace()
  480. tenant = Tenant(name=name)
  481. db.session.add(tenant)
  482. db.session.commit()
  483. tenant.encrypt_public_key = generate_key_pair(tenant.id)
  484. db.session.commit()
  485. return tenant
  486. @staticmethod
  487. def create_owner_tenant_if_not_exist(
  488. account: Account, name: Optional[str] = None, is_setup: Optional[bool] = False
  489. ):
  490. """Check if user have a workspace or not"""
  491. available_ta = (
  492. TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
  493. )
  494. if available_ta:
  495. return
  496. """Create owner tenant if not exist"""
  497. if not FeatureService.get_system_features().is_allow_create_workspace and not is_setup:
  498. raise WorkSpaceNotAllowedCreateError()
  499. if name:
  500. tenant = TenantService.create_tenant(name=name, is_setup=is_setup)
  501. else:
  502. tenant = TenantService.create_tenant(name=f"{account.name}'s Workspace", is_setup=is_setup)
  503. TenantService.create_tenant_member(tenant, account, role="owner")
  504. account.current_tenant = tenant
  505. db.session.commit()
  506. tenant_was_created.send(tenant)
  507. @staticmethod
  508. def create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin:
  509. """Create tenant member"""
  510. if role == TenantAccountRole.OWNER.value:
  511. if TenantService.has_roles(tenant, [TenantAccountRole.OWNER]):
  512. logging.error(f"Tenant {tenant.id} has already an owner.")
  513. raise Exception("Tenant already has an owner.")
  514. ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()
  515. if ta:
  516. ta.role = role
  517. else:
  518. ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
  519. db.session.add(ta)
  520. db.session.commit()
  521. return ta
  522. @staticmethod
  523. def get_join_tenants(account: Account) -> list[Tenant]:
  524. """Get account join tenants"""
  525. return (
  526. db.session.query(Tenant)
  527. .join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
  528. .filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL)
  529. .all()
  530. )
  531. @staticmethod
  532. def get_current_tenant_by_account(account: Account):
  533. """Get tenant by account and add the role"""
  534. tenant = account.current_tenant
  535. if not tenant:
  536. raise TenantNotFoundError("Tenant not found.")
  537. ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
  538. if ta:
  539. tenant.role = ta.role
  540. else:
  541. raise TenantNotFoundError("Tenant not found for the account.")
  542. return tenant
  543. @staticmethod
  544. def switch_tenant(account: Account, tenant_id: Optional[str] = None) -> None:
  545. """Switch the current workspace for the account"""
  546. # Ensure tenant_id is provided
  547. if tenant_id is None:
  548. raise ValueError("Tenant ID must be provided.")
  549. tenant_account_join = (
  550. db.session.query(TenantAccountJoin)
  551. .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
  552. .filter(
  553. TenantAccountJoin.account_id == account.id,
  554. TenantAccountJoin.tenant_id == tenant_id,
  555. Tenant.status == TenantStatus.NORMAL,
  556. )
  557. .first()
  558. )
  559. if not tenant_account_join:
  560. raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.")
  561. else:
  562. TenantAccountJoin.query.filter(
  563. TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id
  564. ).update({"current": False})
  565. tenant_account_join.current = True
  566. # Set the current tenant for the account
  567. account.current_tenant_id = tenant_account_join.tenant_id
  568. db.session.commit()
  569. @staticmethod
  570. def get_tenant_members(tenant: Tenant) -> list[Account]:
  571. """Get tenant members"""
  572. query = (
  573. db.session.query(Account, TenantAccountJoin.role)
  574. .select_from(Account)
  575. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  576. .filter(TenantAccountJoin.tenant_id == tenant.id)
  577. )
  578. # Initialize an empty list to store the updated accounts
  579. updated_accounts = []
  580. for account, role in query:
  581. account.role = role
  582. updated_accounts.append(account)
  583. return updated_accounts
  584. @staticmethod
  585. def get_dataset_operator_members(tenant: Tenant) -> list[Account]:
  586. """Get dataset admin members"""
  587. query = (
  588. db.session.query(Account, TenantAccountJoin.role)
  589. .select_from(Account)
  590. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  591. .filter(TenantAccountJoin.tenant_id == tenant.id)
  592. .filter(TenantAccountJoin.role == "dataset_operator")
  593. )
  594. # Initialize an empty list to store the updated accounts
  595. updated_accounts = []
  596. for account, role in query:
  597. account.role = role
  598. updated_accounts.append(account)
  599. return updated_accounts
  600. @staticmethod
  601. def has_roles(tenant: Tenant, roles: list[TenantAccountRole]) -> bool:
  602. """Check if user has any of the given roles for a tenant"""
  603. if not all(isinstance(role, TenantAccountRole) for role in roles):
  604. raise ValueError("all roles must be TenantAccountRole")
  605. return (
  606. db.session.query(TenantAccountJoin)
  607. .filter(
  608. TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles])
  609. )
  610. .first()
  611. is not None
  612. )
  613. @staticmethod
  614. def get_user_role(account: Account, tenant: Tenant) -> Optional[TenantAccountRole]:
  615. """Get the role of the current account for a given tenant"""
  616. join = (
  617. db.session.query(TenantAccountJoin)
  618. .filter(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id)
  619. .first()
  620. )
  621. return join.role if join else None
  622. @staticmethod
  623. def get_tenant_count() -> int:
  624. """Get tenant count"""
  625. return cast(int, db.session.query(func.count(Tenant.id)).scalar())
  626. @staticmethod
  627. def check_member_permission(tenant: Tenant, operator: Account, member: Account | None, action: str) -> None:
  628. """Check member permission"""
  629. perms = {
  630. "add": [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
  631. "remove": [TenantAccountRole.OWNER],
  632. "update": [TenantAccountRole.OWNER],
  633. }
  634. if action not in {"add", "remove", "update"}:
  635. raise InvalidActionError("Invalid action.")
  636. if member:
  637. if operator.id == member.id:
  638. raise CannotOperateSelfError("Cannot operate self.")
  639. ta_operator = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=operator.id).first()
  640. if not ta_operator or ta_operator.role not in perms[action]:
  641. raise NoPermissionError(f"No permission to {action} member.")
  642. @staticmethod
  643. def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account) -> None:
  644. """Remove member from tenant"""
  645. if operator.id == account.id:
  646. raise CannotOperateSelfError("Cannot operate self.")
  647. TenantService.check_member_permission(tenant, operator, account, "remove")
  648. ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
  649. if not ta:
  650. raise MemberNotInTenantError("Member not in tenant.")
  651. db.session.delete(ta)
  652. db.session.commit()
  653. @staticmethod
  654. def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account) -> None:
  655. """Update member role"""
  656. TenantService.check_member_permission(tenant, operator, member, "update")
  657. target_member_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=member.id).first()
  658. if target_member_join.role == new_role:
  659. raise RoleAlreadyAssignedError("The provided role is already assigned to the member.")
  660. if new_role == "owner":
  661. # Find the current owner and change their role to 'admin'
  662. current_owner_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, role="owner").first()
  663. current_owner_join.role = "admin"
  664. # Update the role of the target member
  665. target_member_join.role = new_role
  666. db.session.commit()
  667. @staticmethod
  668. def dissolve_tenant(tenant: Tenant, operator: Account) -> None:
  669. """Dissolve tenant"""
  670. if not TenantService.check_member_permission(tenant, operator, operator, "remove"):
  671. raise NoPermissionError("No permission to dissolve tenant.")
  672. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id).delete()
  673. db.session.delete(tenant)
  674. db.session.commit()
  675. @staticmethod
  676. def get_custom_config(tenant_id: str) -> dict:
  677. tenant = Tenant.query.filter(Tenant.id == tenant_id).one_or_404()
  678. return cast(dict, tenant.custom_config_dict)
  679. class RegisterService:
  680. @classmethod
  681. def _get_invitation_token_key(cls, token: str) -> str:
  682. return f"member_invite:token:{token}"
  683. @classmethod
  684. def setup(cls, email: str, name: str, password: str, ip_address: str) -> None:
  685. """
  686. Setup dify
  687. :param email: email
  688. :param name: username
  689. :param password: password
  690. :param ip_address: ip address
  691. """
  692. try:
  693. # Register
  694. account = AccountService.create_account(
  695. email=email,
  696. name=name,
  697. interface_language=languages[0],
  698. password=password,
  699. is_setup=True,
  700. )
  701. account.last_login_ip = ip_address
  702. account.initialized_at = datetime.now(UTC).replace(tzinfo=None)
  703. TenantService.create_owner_tenant_if_not_exist(account=account, is_setup=True)
  704. dify_setup = DifySetup(version=dify_config.CURRENT_VERSION)
  705. db.session.add(dify_setup)
  706. db.session.commit()
  707. except Exception as e:
  708. db.session.query(DifySetup).delete()
  709. db.session.query(TenantAccountJoin).delete()
  710. db.session.query(Account).delete()
  711. db.session.query(Tenant).delete()
  712. db.session.commit()
  713. logging.exception(f"Setup account failed, email: {email}, name: {name}")
  714. raise ValueError(f"Setup failed: {e}")
  715. @classmethod
  716. def register(
  717. cls,
  718. email,
  719. name,
  720. password: Optional[str] = None,
  721. open_id: Optional[str] = None,
  722. provider: Optional[str] = None,
  723. language: Optional[str] = None,
  724. status: Optional[AccountStatus] = None,
  725. is_setup: Optional[bool] = False,
  726. create_workspace_required: Optional[bool] = True,
  727. ) -> Account:
  728. db.session.begin_nested()
  729. """Register account"""
  730. try:
  731. account = AccountService.create_account(
  732. email=email,
  733. name=name,
  734. interface_language=language or languages[0],
  735. password=password,
  736. is_setup=is_setup,
  737. )
  738. account.status = AccountStatus.ACTIVE.value if not status else status.value
  739. account.initialized_at = datetime.now(UTC).replace(tzinfo=None)
  740. if open_id is not None and provider is not None:
  741. AccountService.link_account_integrate(provider, open_id, account)
  742. if FeatureService.get_system_features().is_allow_create_workspace and create_workspace_required:
  743. tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
  744. TenantService.create_tenant_member(tenant, account, role="owner")
  745. account.current_tenant = tenant
  746. tenant_was_created.send(tenant)
  747. db.session.commit()
  748. except WorkSpaceNotAllowedCreateError:
  749. db.session.rollback()
  750. except AccountRegisterError as are:
  751. db.session.rollback()
  752. logging.exception("Register failed")
  753. raise are
  754. except Exception as e:
  755. db.session.rollback()
  756. logging.exception("Register failed")
  757. raise AccountRegisterError(f"Registration failed: {e}") from e
  758. return account
  759. @classmethod
  760. def invite_new_member(
  761. cls, tenant: Tenant, email: str, language: str, role: str = "normal", inviter: Account | None = None
  762. ) -> str:
  763. if not inviter:
  764. raise ValueError("Inviter is required")
  765. """Invite new member"""
  766. with Session(db.engine) as session:
  767. account = session.query(Account).filter_by(email=email).first()
  768. if not account:
  769. TenantService.check_member_permission(tenant, inviter, None, "add")
  770. name = email.split("@")[0]
  771. account = cls.register(
  772. email=email, name=name, language=language, status=AccountStatus.PENDING, is_setup=True
  773. )
  774. # Create new tenant member for invited tenant
  775. TenantService.create_tenant_member(tenant, account, role)
  776. TenantService.switch_tenant(account, tenant.id)
  777. else:
  778. TenantService.check_member_permission(tenant, inviter, account, "add")
  779. ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
  780. if not ta:
  781. TenantService.create_tenant_member(tenant, account, role)
  782. # Support resend invitation email when the account is pending status
  783. if account.status != AccountStatus.PENDING.value:
  784. raise AccountAlreadyInTenantError("Account already in tenant.")
  785. token = cls.generate_invite_token(tenant, account)
  786. # send email
  787. send_invite_member_mail_task.delay(
  788. language=account.interface_language,
  789. to=email,
  790. token=token,
  791. inviter_name=inviter.name if inviter else "Dify",
  792. workspace_name=tenant.name,
  793. )
  794. return token
  795. @classmethod
  796. def generate_invite_token(cls, tenant: Tenant, account: Account) -> str:
  797. token = str(uuid.uuid4())
  798. invitation_data = {
  799. "account_id": account.id,
  800. "email": account.email,
  801. "workspace_id": tenant.id,
  802. }
  803. expiry_hours = dify_config.INVITE_EXPIRY_HOURS
  804. redis_client.setex(cls._get_invitation_token_key(token), expiry_hours * 60 * 60, json.dumps(invitation_data))
  805. return token
  806. @classmethod
  807. def is_valid_invite_token(cls, token: str) -> bool:
  808. data = redis_client.get(cls._get_invitation_token_key(token))
  809. return data is not None
  810. @classmethod
  811. def revoke_token(cls, workspace_id: str, email: str, token: str):
  812. if workspace_id and email:
  813. email_hash = sha256(email.encode()).hexdigest()
  814. cache_key = "member_invite_token:{}, {}:{}".format(workspace_id, email_hash, token)
  815. redis_client.delete(cache_key)
  816. else:
  817. redis_client.delete(cls._get_invitation_token_key(token))
  818. @classmethod
  819. def get_invitation_if_token_valid(
  820. cls, workspace_id: Optional[str], email: str, token: str
  821. ) -> Optional[dict[str, Any]]:
  822. invitation_data = cls._get_invitation_by_token(token, workspace_id, email)
  823. if not invitation_data:
  824. return None
  825. tenant = (
  826. db.session.query(Tenant)
  827. .filter(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal")
  828. .first()
  829. )
  830. if not tenant:
  831. return None
  832. tenant_account = (
  833. db.session.query(Account, TenantAccountJoin.role)
  834. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  835. .filter(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id)
  836. .first()
  837. )
  838. if not tenant_account:
  839. return None
  840. account = tenant_account[0]
  841. if not account:
  842. return None
  843. if invitation_data["account_id"] != str(account.id):
  844. return None
  845. return {
  846. "account": account,
  847. "data": invitation_data,
  848. "tenant": tenant,
  849. }
  850. @classmethod
  851. def _get_invitation_by_token(
  852. cls, token: str, workspace_id: Optional[str] = None, email: Optional[str] = None
  853. ) -> Optional[dict[str, str]]:
  854. if workspace_id is not None and email is not None:
  855. email_hash = sha256(email.encode()).hexdigest()
  856. cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
  857. account_id = redis_client.get(cache_key)
  858. if not account_id:
  859. return None
  860. return {
  861. "account_id": account_id.decode("utf-8"),
  862. "email": email,
  863. "workspace_id": workspace_id,
  864. }
  865. else:
  866. data = redis_client.get(cls._get_invitation_token_key(token))
  867. if not data:
  868. return None
  869. invitation: dict = json.loads(data)
  870. return invitation
  871. def _generate_refresh_token(length: int = 64):
  872. token = secrets.token_hex(length)
  873. return token