account_service.py 35 KB

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