open-webui/backend/open_webui/utils/auth.py

236 lines
6.6 KiB
Python
Raw Normal View History

2024-08-27 22:10:27 +00:00
import logging
import uuid
2024-11-15 09:29:07 +00:00
import jwt
2025-02-16 08:11:18 +00:00
import base64
import hmac
import hashlib
2025-02-17 02:35:09 +00:00
import requests
2025-02-18 05:34:06 +00:00
import os
2025-02-17 02:35:09 +00:00
2024-11-15 09:29:07 +00:00
2024-08-27 22:10:27 +00:00
from datetime import UTC, datetime, timedelta
2024-11-15 09:29:07 +00:00
from typing import Optional, Union, List, Dict
2024-12-10 08:54:13 +00:00
from open_webui.models.users import Users
2024-11-15 09:29:07 +00:00
from open_webui.constants import ERROR_MESSAGES
2025-02-18 05:34:06 +00:00
from open_webui.env import WEBUI_SECRET_KEY, TRUSTED_SIGNATURE_KEY, STATIC_DIR
2024-11-15 09:29:07 +00:00
2024-10-15 09:48:41 +00:00
from fastapi import Depends, HTTPException, Request, Response, status
2024-08-27 22:10:27 +00:00
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from passlib.context import CryptContext
2023-11-19 00:47:12 +00:00
2024-01-05 20:22:27 +00:00
logging.getLogger("passlib").setLevel(logging.ERROR)
2024-08-25 14:52:36 +00:00
SESSION_SECRET = WEBUI_SECRET_KEY
2023-11-19 00:47:12 +00:00
ALGORITHM = "HS256"
##############
# Auth Utils
##############
2025-02-16 08:11:18 +00:00
def verify_signature(payload: str, signature: str) -> bool:
"""
Verifies the HMAC signature of the received payload.
"""
try:
expected_signature = base64.b64encode(
hmac.new(TRUSTED_SIGNATURE_KEY, payload.encode(), hashlib.sha256).digest()
).decode()
# Compare securely to prevent timing attacks
return hmac.compare_digest(expected_signature, signature)
except Exception:
return False
2025-02-17 02:35:09 +00:00
2025-02-18 05:34:06 +00:00
def override_static(path: str, content: str):
# Ensure path is safe
if "/" in path or ".." in path:
print(f"Invalid path: {path}")
return
file_path = os.path.join(STATIC_DIR, path)
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, "wb") as f:
f.write(base64.b64decode(content)) # Convert Base64 back to raw binary
2025-02-17 02:35:09 +00:00
def get_license_data(app, key):
if key:
try:
res = requests.post(
"https://api.openwebui.com/api/v1/license",
json={"key": key, "version": "1"},
timeout=5,
)
if getattr(res, "ok", False):
payload = getattr(res, "json", lambda: {})()
for k, v in payload.items():
if k == "resources":
for p, c in v.items():
globals().get("override_static", lambda a, b: None)(p, c)
elif k == "user_count":
setattr(app.state, "USER_COUNT", v)
elif k == "webui_name":
setattr(app.state, "WEBUI_NAME", v)
return True
else:
print(
f"License: retrieval issue: {getattr(res, 'text', 'unknown error')}"
)
except Exception as ex:
print(f"License: Uncaught Exception: {ex}")
return False
2023-11-19 00:47:12 +00:00
2024-06-19 21:38:09 +00:00
bearer_security = HTTPBearer(auto_error=False)
2023-11-19 00:47:12 +00:00
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
def verify_password(plain_password, hashed_password):
2024-01-05 20:22:27 +00:00
return (
pwd_context.verify(plain_password, hashed_password) if hashed_password else None
)
2023-11-19 00:47:12 +00:00
def get_password_hash(password):
return pwd_context.hash(password)
2024-01-05 20:22:27 +00:00
def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
2023-11-19 00:47:12 +00:00
payload = data.copy()
if expires_delta:
expire = datetime.now(UTC) + expires_delta
2023-11-19 00:47:12 +00:00
payload.update({"exp": expire})
encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)
2023-11-19 00:47:12 +00:00
return encoded_jwt
def decode_token(token: str) -> Optional[dict]:
try:
decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])
2023-11-19 00:47:12 +00:00
return decoded
2024-08-03 13:24:26 +00:00
except Exception:
2023-11-19 00:47:12 +00:00
return None
def extract_token_from_auth_header(auth_header: str):
2024-01-05 20:22:27 +00:00
return auth_header[len("Bearer ") :]
2023-11-19 00:47:12 +00:00
2024-03-26 10:22:17 +00:00
def create_api_key():
key = str(uuid.uuid4()).replace("-", "")
return f"sk-{key}"
2024-02-24 06:44:56 +00:00
def get_http_authorization_cred(auth_header: str):
try:
scheme, credentials = auth_header.split(" ")
2024-02-25 06:10:43 +00:00
return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
2024-08-03 13:24:26 +00:00
except Exception:
2024-02-24 06:44:56 +00:00
raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)
2024-02-11 01:54:33 +00:00
def get_current_user(
2024-06-19 21:38:09 +00:00
request: Request,
2024-02-11 01:54:33 +00:00
auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
):
2024-06-19 21:49:35 +00:00
token = None
2024-06-19 21:38:09 +00:00
if auth_token is not None:
token = auth_token.credentials
2024-06-19 21:49:35 +00:00
if token is None and "token" in request.cookies:
token = request.cookies.get("token")
if token is None:
raise HTTPException(status_code=403, detail="Not authenticated")
2024-03-26 10:22:17 +00:00
# auth by api key
2024-06-19 21:38:09 +00:00
if token.startswith("sk-"):
2024-11-19 20:17:23 +00:00
if not request.state.enable_api_key:
raise HTTPException(
status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
)
2024-12-25 06:32:34 +00:00
2024-12-27 08:32:25 +00:00
if request.app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS:
2024-12-27 04:58:46 +00:00
allowed_paths = [
path.strip()
2025-01-03 21:08:21 +00:00
for path in str(
request.app.state.config.API_KEY_ALLOWED_ENDPOINTS
).split(",")
2024-12-27 04:58:46 +00:00
]
if request.url.path not in allowed_paths:
raise HTTPException(
status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
)
2024-12-25 06:32:34 +00:00
return get_current_user_by_api_key(token)
2024-06-19 21:38:09 +00:00
2024-03-26 10:22:17 +00:00
# auth by jwt token
2024-11-06 05:14:02 +00:00
try:
data = decode_token(token)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid token",
)
2024-08-03 13:24:26 +00:00
if data is not None and "id" in data:
user = Users.get_user_by_id(data["id"])
if user is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
2023-11-19 00:47:12 +00:00
)
2024-04-27 23:38:51 +00:00
else:
Users.update_user_last_active_by_id(user.id)
return user
else:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.UNAUTHORIZED,
)
2024-04-02 16:42:45 +00:00
2024-06-24 11:45:33 +00:00
def get_current_user_by_api_key(api_key: str):
user = Users.get_user_by_api_key(api_key)
2024-04-27 23:38:51 +00:00
2024-03-26 10:22:17 +00:00
if user is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
)
2024-04-27 23:38:51 +00:00
else:
2024-06-24 11:45:33 +00:00
Users.update_user_last_active_by_id(user.id)
2024-04-27 23:38:51 +00:00
2024-03-26 10:22:17 +00:00
return user
2024-04-02 16:42:45 +00:00
2024-02-11 01:54:33 +00:00
def get_verified_user(user=Depends(get_current_user)):
if user.role not in {"user", "admin"}:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
2024-02-11 01:54:33 +00:00
return user
2024-02-11 01:54:33 +00:00
def get_admin_user(user=Depends(get_current_user)):
if user.role != "admin":
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
2024-02-11 01:54:33 +00:00
return user