utils.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. import logging
  2. import uuid
  3. import jwt
  4. from datetime import UTC, datetime, timedelta
  5. from typing import Optional, Union, List, Dict
  6. from open_webui.apps.webui.models.users import Users
  7. from open_webui.apps.webui.models.groups import Groups
  8. from open_webui.constants import ERROR_MESSAGES
  9. from open_webui.env import WEBUI_SECRET_KEY
  10. from fastapi import Depends, HTTPException, Request, Response, status
  11. from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
  12. from passlib.context import CryptContext
  13. logging.getLogger("passlib").setLevel(logging.ERROR)
  14. SESSION_SECRET = WEBUI_SECRET_KEY
  15. ALGORITHM = "HS256"
  16. ##############
  17. # Auth Utils
  18. ##############
  19. bearer_security = HTTPBearer(auto_error=False)
  20. pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
  21. def verify_password(plain_password, hashed_password):
  22. return (
  23. pwd_context.verify(plain_password, hashed_password) if hashed_password else None
  24. )
  25. def get_password_hash(password):
  26. return pwd_context.hash(password)
  27. def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
  28. payload = data.copy()
  29. if expires_delta:
  30. expire = datetime.now(UTC) + expires_delta
  31. payload.update({"exp": expire})
  32. encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)
  33. return encoded_jwt
  34. def decode_token(token: str) -> Optional[dict]:
  35. try:
  36. decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])
  37. return decoded
  38. except Exception:
  39. return None
  40. def extract_token_from_auth_header(auth_header: str):
  41. return auth_header[len("Bearer ") :]
  42. def create_api_key():
  43. key = str(uuid.uuid4()).replace("-", "")
  44. return f"sk-{key}"
  45. def get_http_authorization_cred(auth_header: str):
  46. try:
  47. scheme, credentials = auth_header.split(" ")
  48. return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
  49. except Exception:
  50. raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)
  51. def get_current_user(
  52. request: Request,
  53. auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
  54. ):
  55. token = None
  56. if auth_token is not None:
  57. token = auth_token.credentials
  58. if token is None and "token" in request.cookies:
  59. token = request.cookies.get("token")
  60. if token is None:
  61. raise HTTPException(status_code=403, detail="Not authenticated")
  62. # auth by api key
  63. if token.startswith("sk-"):
  64. return get_current_user_by_api_key(token)
  65. # auth by jwt token
  66. try:
  67. data = decode_token(token)
  68. except Exception as e:
  69. raise HTTPException(
  70. status_code=status.HTTP_401_UNAUTHORIZED,
  71. detail="Invalid token",
  72. )
  73. if data is not None and "id" in data:
  74. user = Users.get_user_by_id(data["id"])
  75. if user is None:
  76. raise HTTPException(
  77. status_code=status.HTTP_401_UNAUTHORIZED,
  78. detail=ERROR_MESSAGES.INVALID_TOKEN,
  79. )
  80. else:
  81. Users.update_user_last_active_by_id(user.id)
  82. return user
  83. else:
  84. raise HTTPException(
  85. status_code=status.HTTP_401_UNAUTHORIZED,
  86. detail=ERROR_MESSAGES.UNAUTHORIZED,
  87. )
  88. def get_current_user_by_api_key(api_key: str):
  89. user = Users.get_user_by_api_key(api_key)
  90. if user is None:
  91. raise HTTPException(
  92. status_code=status.HTTP_401_UNAUTHORIZED,
  93. detail=ERROR_MESSAGES.INVALID_TOKEN,
  94. )
  95. else:
  96. Users.update_user_last_active_by_id(user.id)
  97. return user
  98. def get_verified_user(user=Depends(get_current_user)):
  99. if user.role not in {"user", "admin"}:
  100. raise HTTPException(
  101. status_code=status.HTTP_401_UNAUTHORIZED,
  102. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  103. )
  104. return user
  105. def get_admin_user(user=Depends(get_current_user)):
  106. if user.role != "admin":
  107. raise HTTPException(
  108. status_code=status.HTTP_401_UNAUTHORIZED,
  109. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  110. )
  111. return user
  112. def has_permission(
  113. user_id: str,
  114. permission_key: str,
  115. default_permissions: Dict[str, bool] = {},
  116. ) -> bool:
  117. """
  118. Check if a user has a specific permission by checking the group permissions
  119. and falls back to default permissions if not found in any group.
  120. Permission keys can be hierarchical and separated by dots ('.').
  121. """
  122. def get_permission(permissions: Dict[str, bool], keys: List[str]) -> bool:
  123. """Traverse permissions dict using a list of keys (from dot-split permission_key)."""
  124. for key in keys:
  125. if key not in permissions:
  126. return False # If any part of the hierarchy is missing, deny access
  127. permissions = permissions[key] # Go one level deeper
  128. return bool(permissions) # Return the boolean at the final level
  129. permission_hierarchy = permission_key.split(".")
  130. # Retrieve user group permissions
  131. user_groups = Groups.get_groups_by_member_id(user_id)
  132. for group in user_groups:
  133. group_permissions = group.permissions
  134. if get_permission(group_permissions, permission_hierarchy):
  135. return True
  136. # Check default permissions afterwards if the group permissions don't allow it
  137. return get_permission(default_permissions, permission_hierarchy)
  138. def has_access(
  139. user_id: str,
  140. action: str = "write",
  141. access_control: Optional[dict] = None,
  142. ) -> bool:
  143. if access_control is None:
  144. return action == "read"
  145. user_groups = Groups.get_groups_by_member_id(user_id)
  146. user_group_ids = [group.id for group in user_groups]
  147. permission_access = access_control.get(action, {})
  148. permitted_group_ids = permission_access.get("group_ids", [])
  149. permitted_user_ids = permission_access.get("user_ids", [])
  150. return user_id in permitted_user_ids or any(
  151. group_id in permitted_group_ids for group_id in user_group_ids
  152. )