auth.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235
  1. import logging
  2. import uuid
  3. import jwt
  4. import base64
  5. import hmac
  6. import hashlib
  7. import requests
  8. import os
  9. from datetime import UTC, datetime, timedelta
  10. from typing import Optional, Union, List, Dict
  11. from open_webui.models.users import Users
  12. from open_webui.constants import ERROR_MESSAGES
  13. from open_webui.env import WEBUI_SECRET_KEY, TRUSTED_SIGNATURE_KEY, STATIC_DIR
  14. from fastapi import Depends, HTTPException, Request, Response, status
  15. from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
  16. from passlib.context import CryptContext
  17. logging.getLogger("passlib").setLevel(logging.ERROR)
  18. SESSION_SECRET = WEBUI_SECRET_KEY
  19. ALGORITHM = "HS256"
  20. ##############
  21. # Auth Utils
  22. ##############
  23. def verify_signature(payload: str, signature: str) -> bool:
  24. """
  25. Verifies the HMAC signature of the received payload.
  26. """
  27. try:
  28. expected_signature = base64.b64encode(
  29. hmac.new(TRUSTED_SIGNATURE_KEY, payload.encode(), hashlib.sha256).digest()
  30. ).decode()
  31. # Compare securely to prevent timing attacks
  32. return hmac.compare_digest(expected_signature, signature)
  33. except Exception:
  34. return False
  35. def override_static(path: str, content: str):
  36. # Ensure path is safe
  37. if "/" in path or ".." in path:
  38. print(f"Invalid path: {path}")
  39. return
  40. file_path = os.path.join(STATIC_DIR, path)
  41. os.makedirs(os.path.dirname(file_path), exist_ok=True)
  42. with open(file_path, "wb") as f:
  43. f.write(base64.b64decode(content)) # Convert Base64 back to raw binary
  44. def get_license_data(app, key):
  45. if key:
  46. try:
  47. res = requests.post(
  48. "https://api.openwebui.com/api/v1/license",
  49. json={"key": key, "version": "1"},
  50. timeout=5,
  51. )
  52. if getattr(res, "ok", False):
  53. payload = getattr(res, "json", lambda: {})()
  54. for k, v in payload.items():
  55. if k == "resources":
  56. for p, c in v.items():
  57. globals().get("override_static", lambda a, b: None)(p, c)
  58. elif k == "user_count":
  59. setattr(app.state, "USER_COUNT", v)
  60. elif k == "webui_name":
  61. setattr(app.state, "WEBUI_NAME", v)
  62. return True
  63. else:
  64. print(
  65. f"License: retrieval issue: {getattr(res, 'text', 'unknown error')}"
  66. )
  67. except Exception as ex:
  68. print(f"License: Uncaught Exception: {ex}")
  69. return False
  70. bearer_security = HTTPBearer(auto_error=False)
  71. pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
  72. def verify_password(plain_password, hashed_password):
  73. return (
  74. pwd_context.verify(plain_password, hashed_password) if hashed_password else None
  75. )
  76. def get_password_hash(password):
  77. return pwd_context.hash(password)
  78. def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
  79. payload = data.copy()
  80. if expires_delta:
  81. expire = datetime.now(UTC) + expires_delta
  82. payload.update({"exp": expire})
  83. encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)
  84. return encoded_jwt
  85. def decode_token(token: str) -> Optional[dict]:
  86. try:
  87. decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])
  88. return decoded
  89. except Exception:
  90. return None
  91. def extract_token_from_auth_header(auth_header: str):
  92. return auth_header[len("Bearer ") :]
  93. def create_api_key():
  94. key = str(uuid.uuid4()).replace("-", "")
  95. return f"sk-{key}"
  96. def get_http_authorization_cred(auth_header: str):
  97. try:
  98. scheme, credentials = auth_header.split(" ")
  99. return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
  100. except Exception:
  101. raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)
  102. def get_current_user(
  103. request: Request,
  104. auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
  105. ):
  106. token = None
  107. if auth_token is not None:
  108. token = auth_token.credentials
  109. if token is None and "token" in request.cookies:
  110. token = request.cookies.get("token")
  111. if token is None:
  112. raise HTTPException(status_code=403, detail="Not authenticated")
  113. # auth by api key
  114. if token.startswith("sk-"):
  115. if not request.state.enable_api_key:
  116. raise HTTPException(
  117. status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
  118. )
  119. if request.app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS:
  120. allowed_paths = [
  121. path.strip()
  122. for path in str(
  123. request.app.state.config.API_KEY_ALLOWED_ENDPOINTS
  124. ).split(",")
  125. ]
  126. if request.url.path not in allowed_paths:
  127. raise HTTPException(
  128. status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
  129. )
  130. return get_current_user_by_api_key(token)
  131. # auth by jwt token
  132. try:
  133. data = decode_token(token)
  134. except Exception as e:
  135. raise HTTPException(
  136. status_code=status.HTTP_401_UNAUTHORIZED,
  137. detail="Invalid token",
  138. )
  139. if data is not None and "id" in data:
  140. user = Users.get_user_by_id(data["id"])
  141. if user is None:
  142. raise HTTPException(
  143. status_code=status.HTTP_401_UNAUTHORIZED,
  144. detail=ERROR_MESSAGES.INVALID_TOKEN,
  145. )
  146. else:
  147. Users.update_user_last_active_by_id(user.id)
  148. return user
  149. else:
  150. raise HTTPException(
  151. status_code=status.HTTP_401_UNAUTHORIZED,
  152. detail=ERROR_MESSAGES.UNAUTHORIZED,
  153. )
  154. def get_current_user_by_api_key(api_key: str):
  155. user = Users.get_user_by_api_key(api_key)
  156. if user is None:
  157. raise HTTPException(
  158. status_code=status.HTTP_401_UNAUTHORIZED,
  159. detail=ERROR_MESSAGES.INVALID_TOKEN,
  160. )
  161. else:
  162. Users.update_user_last_active_by_id(user.id)
  163. return user
  164. def get_verified_user(user=Depends(get_current_user)):
  165. if user.role not in {"user", "admin"}:
  166. raise HTTPException(
  167. status_code=status.HTTP_401_UNAUTHORIZED,
  168. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  169. )
  170. return user
  171. def get_admin_user(user=Depends(get_current_user)):
  172. if user.role != "admin":
  173. raise HTTPException(
  174. status_code=status.HTTP_401_UNAUTHORIZED,
  175. detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
  176. )
  177. return user