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

import logging
import time
from string import punctuation
from typing import Any, Optional
from urllib.parse import quote

import jwt
from fastapi import Request
from open_webui.env import (
    ENABLE_FORWARD_USER_INFO_HEADERS,
    FORWARD_SESSION_INFO_HEADER_CHAT_ID,
    FORWARD_USER_INFO_HEADER_AUTH_TYPE,
    FORWARD_USER_INFO_HEADER_JWT,
    FORWARD_USER_INFO_HEADER_JWT_EXPIRES_SECONDS,
    FORWARD_USER_INFO_HEADER_JWT_SECRET,
    FORWARD_USER_INFO_HEADER_USER_EMAIL,
    FORWARD_USER_INFO_HEADER_USER_ID,
    FORWARD_USER_INFO_HEADER_USER_NAME,
    FORWARD_USER_INFO_HEADER_USER_ROLE,
)
from open_webui.models.groups import Groups

log = logging.getLogger(__name__)

USER_GROUPS_PLACEHOLDERS = ('{{USER_GROUPS}}', '{{USER_GROUP_IDS}}')


def normalize_bearer_token(token: Any) -> str:
    return token.strip() if isinstance(token, str) else token or ''


def bearer_auth_header(token: Any) -> dict[str, str]:
    token = normalize_bearer_token(token)
    return {'Authorization': f'Bearer {token}'} if token else {}


def get_json_bearer_headers(token: Any = '') -> dict[str, str]:
    return {'Content-Type': 'application/json', **bearer_auth_header(token)}


def _mint_forward_user_jwt(user: Any) -> str:
    now = int(time.time())
    payload = {
        'sub': str(user.id),
        'email': str(user.email),
        'name': str(user.name),
        'role': str(user.role),
        'iss': 'open-webui',
        'iat': now,
        'exp': now + FORWARD_USER_INFO_HEADER_JWT_EXPIRES_SECONDS,
    }
    return jwt.encode(payload, FORWARD_USER_INFO_HEADER_JWT_SECRET, algorithm='HS256')


def include_user_info_headers(headers: dict, user: Optional[Any] = None, *, request=None) -> dict:
    """
    Forward user identity to external backends: signed JWT in
    FORWARD_USER_INFO_HEADER_JWT if FORWARD_USER_INFO_HEADER_JWT_SECRET is set;
    otherwise the legacy X-OpenWebUI-User-* headers.
    Include the verified incoming auth type when a request provides it.
    """
    if user is None:
        return headers

    auth_type = getattr(getattr(request, 'state', None), 'auth_type', None)
    if auth_type in ('api_key', 'jwt'):
        headers = {**headers, FORWARD_USER_INFO_HEADER_AUTH_TYPE: auth_type}

    if FORWARD_USER_INFO_HEADER_JWT_SECRET:
        try:
            token = _mint_forward_user_jwt(user)
            return {**headers, FORWARD_USER_INFO_HEADER_JWT: token}
        except Exception:
            log.exception(
                'Failed to mint %s; falling back to plain user-info headers.',
                FORWARD_USER_INFO_HEADER_JWT,
            )

    return {
        **headers,
        FORWARD_USER_INFO_HEADER_USER_NAME: quote(user.name.strip(), safe=' '),
        FORWARD_USER_INFO_HEADER_USER_ID: user.id,
        FORWARD_USER_INFO_HEADER_USER_EMAIL: user.email.strip(),
        FORWARD_USER_INFO_HEADER_USER_ROLE: user.role,
    }


def custom_headers_require_user_groups(custom_headers: Optional[dict]) -> bool:
    if not custom_headers or not isinstance(custom_headers, dict):
        return False
    return any(
        placeholder in str(value) for value in custom_headers.values() for placeholder in USER_GROUPS_PLACEHOLDERS
    )


async def get_user_groups_for_custom_headers(
    custom_headers: Optional[dict], user: Optional[Any] = None
) -> Optional[list]:
    """Fetch the user's groups only when a header value actually references a groups placeholder."""
    if user is None or not custom_headers_require_user_groups(custom_headers):
        return None

    try:
        return await Groups.get_groups_by_member_id(user.id, include_inherited=True)
    except Exception:
        log.exception('Failed to resolve user groups for custom headers')
        return None


async def get_custom_headers(custom_headers: dict, user=None, metadata: dict = None, request=None) -> dict:
    user_groups = await get_user_groups_for_custom_headers(custom_headers, user)
    return parse_custom_headers(custom_headers, user, metadata, request=request, user_groups=user_groups)


def parse_custom_headers(
    custom_headers: dict, user=None, metadata: dict = None, request=None, user_groups: Optional[list] = None
) -> dict:
    if not custom_headers or not isinstance(custom_headers, dict):
        return {}

    metadata = metadata or {}

    # UA from the live request; fall back to metadata for detached RAG/tool calls.
    user_agent = ''
    if request is not None:
        try:
            user_agent = request.headers.get('user-agent', '') or ''
        except Exception:
            user_agent = ''
    if not user_agent:
        user_agent = metadata.get('user_agent', '') or ''

    # Extract user_message info for tree mapping
    user_message = metadata.get('user_message') or {}
    user_message_id = metadata.get('user_message_id', '') or (user_message.get('id', '') if user_message else '')
    user_message_parent_id = user_message.get('parentId', '') if user_message else ''

    template_vars = {
        '{{CHAT_ID}}': metadata.get('chat_id', '') or '',
        '{{MESSAGE_ID}}': metadata.get('message_id', '') or '',
        '{{USER_MESSAGE_ID}}': user_message_id or '',
        '{{USER_MESSAGE_PARENT_ID}}': user_message_parent_id or '',
        '{{FILE_ID}}': metadata.get('file_id', '') or '',
        '{{FILE_NAME}}': metadata.get('file_name', '') or '',
        '{{FILE_CONTENT_TYPE}}': metadata.get('file_content_type', '') or '',
        '{{TASK}}': metadata.get('task', '') or '',
        '{{MODEL_ID}}': metadata.get('model_id', '') or '',
        '{{USER_ID}}': (user.id if user else '') or '',
        '{{USER_NAME}}': (user.name.strip() if user else '') or '',
        '{{USER_EMAIL}}': (user.email.strip() if user else '') or '',
        '{{USER_ROLE}}': (user.role if user else '') or '',
        '{{USER_GROUPS}}': ','.join(group.name.strip() for group in user_groups) if user_groups else '',
        '{{USER_GROUP_IDS}}': ','.join(group.id for group in user_groups) if user_groups else '',
        '{{USER_AGENT}}': user_agent,
        '{{AUTH_TYPE}}': getattr(getattr(request, 'state', None), 'auth_type', None) or '',
    }

    parsed_headers = {}
    for key, value in custom_headers.items():
        if not isinstance(value, str):
            value = str(value)
        for token, val in template_vars.items():
            value = value.replace(token, val)
        # Encode Unicode and controls after substitution; preserve ASCII header syntax and existing escapes.
        parsed_headers[key] = quote(value, safe=punctuation + ' \t')

    return parsed_headers


async def get_headers_and_cookies(
    request: Request,
    url,
    key=None,
    config=None,
    metadata: dict | None = None,
    user=None,
):
    config = config or {}
    cookies = getattr(request, 'cookies', {}) if config.get('forward_cookies', False) else {}
    headers = {
        'Content-Type': 'application/json',
        **(
            {
                # LICENSE covers this Open WebUI upstream metadata identifier.
                # Do not alter, remove, obscure, or replace it except as LICENSE permits:
                # https://docs.openwebui.com/license.
                'HTTP-Referer': 'https://openwebui.com/',
                'X-Title': 'Open WebUI',
            }
            if 'openrouter.ai' in url
            else {}
        ),
    }

    if ENABLE_FORWARD_USER_INFO_HEADERS and user:
        headers = include_user_info_headers(headers, user, request=request)
        if metadata and metadata.get('chat_id'):
            headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = metadata.get('chat_id')

    token = None
    auth_type = config.get('auth_type')

    if auth_type == 'bearer' or auth_type is None:
        # Default to bearer if not specified
        token = key
    elif auth_type == 'none':
        token = None
    elif auth_type == 'session':
        token = request.state.token.credentials
    elif auth_type == 'system_oauth':
        oauth_token = None
        try:
            if request.cookies.get('oauth_session_id', None):
                oauth_token = await request.app.state.oauth_manager.get_oauth_token(
                    user.id,
                    request.cookies.get('oauth_session_id', None),
                )
        except Exception as e:
            log.error(f'Error getting OAuth token: {e}')

        if oauth_token:
            token = f'{oauth_token.get("access_token", "")}'

    elif auth_type in ('azure_ad', 'microsoft_entra_id'):
        token = get_microsoft_entra_id_access_token()

    if token:
        headers['Authorization'] = f'Bearer {token}'

    if config.get('headers') and isinstance(config.get('headers'), dict):
        custom_headers = await get_custom_headers(config.get('headers'), user, metadata, request=request)
        headers.update(custom_headers)

    return headers, cookies


def get_microsoft_entra_id_access_token():
    """
    Get Microsoft Entra ID access token using DefaultAzureCredential for Azure OpenAI.
    Returns the token string or None if authentication fails.
    """
    from azure.identity import DefaultAzureCredential, get_bearer_token_provider

    try:
        token_provider = get_bearer_token_provider(
            DefaultAzureCredential(), 'https://cognitiveservices.azure.com/.default'
        )
        return token_provider()
    except Exception as e:
        log.error(f'Error getting Microsoft Entra ID access token: {e}')
        return None