auth.py 8.5 KB

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