dify-1.8.1/api/core/file/helpers.py

import base64
import hashlib
import hmac
import os
import time

from configs import dify_config


def get_signed_file_url(upload_file_id: str) -> str:
    url = f"{dify_config.FILES_URL}/files/{upload_file_id}/file-preview"

    timestamp = str(int(time.time()))
    nonce = os.urandom(16).hex()
    key = dify_config.SECRET_KEY.encode()
    msg = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
    sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
    encoded_sign = base64.urlsafe_b64encode(sign).decode()

    return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"


def get_signed_file_url_for_plugin(filename: str, mimetype: str, tenant_id: str, user_id: str) -> str:
    # Plugin access should use internal URL for Docker network communication
    base_url = dify_config.INTERNAL_FILES_URL or dify_config.FILES_URL
    url = f"{base_url}/files/upload/for-plugin"

    if user_id is None:
        user_id = "DEFAULT-USER"

    timestamp = str(int(time.time()))
    nonce = os.urandom(16).hex()
    key = dify_config.SECRET_KEY.encode()
    msg = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}"
    sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
    encoded_sign = base64.urlsafe_b64encode(sign).decode()

    return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}&user_id={user_id}&tenant_id={tenant_id}"


def verify_plugin_file_signature(
    *, filename: str, mimetype: str, tenant_id: str, user_id: str | None, timestamp: str, nonce: str, sign: str
) -> bool:
    if user_id is None:
        user_id = "DEFAULT-USER"

    data_to_sign = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}"
    secret_key = dify_config.SECRET_KEY.encode()
    recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
    recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()

    # verify signature
    if sign != recalculated_encoded_sign:
        return False

    current_time = int(time.time())
    return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT


def verify_image_signature(*, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
    data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
    secret_key = dify_config.SECRET_KEY.encode()
    recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
    recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()

    # verify signature
    if sign != recalculated_encoded_sign:
        return False

    current_time = int(time.time())
    return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT


def verify_file_signature(*, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
    data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
    secret_key = dify_config.SECRET_KEY.encode()
    recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
    recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()

    # verify signature
    if sign != recalculated_encoded_sign:
        return False

    current_time = int(time.time())
    return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT


def get_signed_document_proxy_url(dataset_id: str, document_id: str, filename: str) -> str:
    """Generate signed URL for document file proxy (for KKFileView)"""
    from urllib.parse import quote

    base_url = dify_config.CONSOLE_API_URL or ""
    # Encode filename for URL path
    encoded_filename = quote(filename)
    url = f"{base_url}/console/api/datasets/{dataset_id}/documents/{document_id}/file-proxy/{encoded_filename}"

    timestamp = str(int(time.time()))
    nonce = os.urandom(16).hex()
    key = dify_config.SECRET_KEY.encode()
    msg = f"document-proxy|{dataset_id}|{document_id}|{filename}|{timestamp}|{nonce}"
    sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
    encoded_sign = base64.urlsafe_b64encode(sign).decode()

    return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"


def verify_document_proxy_signature(
    *, dataset_id: str, document_id: str, filename: str, timestamp: str, nonce: str, sign: str
) -> bool:
    """Verify signature for document file proxy"""
    data_to_sign = f"document-proxy|{dataset_id}|{document_id}|{filename}|{timestamp}|{nonce}"
    secret_key = dify_config.SECRET_KEY.encode()
    recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
    recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()

    # verify signature
    if sign != recalculated_encoded_sign:
        return False

    current_time = int(time.time())
    return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT


def generate_document_proxy_token(dataset_id: str, document_id: str, filename: str) -> str:
    """Generate a secure token for document file proxy access

    The token encodes all necessary information (dataset_id, document_id, filename,
    timestamp, nonce, and signature) into a single URL-safe base64 string.

    Args:
        dataset_id: Dataset UUID
        document_id: Document UUID
        filename: Original filename

    Returns:
        URL-safe base64-encoded token
    """
    import json

    timestamp = str(int(time.time()))
    nonce = os.urandom(16).hex()
    key = dify_config.SECRET_KEY.encode()
    msg = f"document-proxy|{dataset_id}|{document_id}|{filename}|{timestamp}|{nonce}"
    sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
    encoded_sign = base64.urlsafe_b64encode(sign).decode()

    # Create payload with all parameters
    payload = {
        "dataset_id": dataset_id,
        "document_id": document_id,
        "filename": filename,
        "timestamp": timestamp,
        "nonce": nonce,
        "sign": encoded_sign,
    }

    # Encode payload as base64 URL-safe string
    token_bytes = json.dumps(payload).encode('utf-8')
    token = base64.urlsafe_b64encode(token_bytes).decode('utf-8').rstrip('=')

    return token


def verify_document_proxy_token(token: str) -> dict[str, str] | None:
    """Verify and decode a document file proxy token

    Args:
        token: Base64-encoded token from URL

    Returns:
        Dict with verified parameters (dataset_id, document_id, filename) if valid,
        None otherwise
    """
    import json

    try:
        # Add padding if needed
        padding = 4 - (len(token) % 4)
        if padding != 4:
            token += '=' * padding

        # Decode base64 token
        token_bytes = base64.urlsafe_b64decode(token)
        payload = json.loads(token_bytes.decode('utf-8'))

        # Extract parameters
        dataset_id = payload.get("dataset_id")
        document_id = payload.get("document_id")
        filename = payload.get("filename")
        timestamp = payload.get("timestamp")
        nonce = payload.get("nonce")
        sign = payload.get("sign")

        if not all([dataset_id, document_id, filename, timestamp, nonce, sign]):
            return None

        # Verify signature using existing function
        if not verify_document_proxy_signature(
            dataset_id=dataset_id,
            document_id=document_id,
            filename=filename,
            timestamp=timestamp,
            nonce=nonce,
            sign=sign
        ):
            return None

        # Return verified parameters
        return {
            "dataset_id": dataset_id,
            "document_id": document_id,
            "filename": filename,
        }

    except (ValueError, AttributeError, KeyError, json.JSONDecodeError):
        return None