dify-1.8.1/api/services/document_download_service.py

"""
Document Download Permission and Statistics Service

This service provides simple document download permission control and statistics tracking using Redis.
All data is stored in Redis with automatic expiration.
"""
import json
import logging
import time
from datetime import datetime
from typing import Any, Optional

from extensions.ext_redis import redis_client, redis_fallback
from models.account import Account

logger = logging.getLogger(__name__)


class DocumentDownloadService:
    """Document download permission and statistics service - simplified version"""

    # Redis Key prefixes
    PERMISSION_PREFIX = "doc:perm"  # Permission configuration
    LOG_PREFIX = "doc:log"  # Download logs
    COUNT_PREFIX = "doc:count"  # Download count per user

    # Data retention time (90 days)
    RETENTION_DAYS = 90
    RETENTION_SECONDS = RETENTION_DAYS * 24 * 60 * 60

    # Maximum log entries per document
    MAX_LOG_ENTRIES = 1000

    # ========== Permission Management ==========

    @classmethod
    def _get_permission_key(cls, tenant_id: str, document_id: str) -> str:
        """Get Redis key for permission configuration."""
        return f"{cls.PERMISSION_PREFIX}:{tenant_id}:{document_id}"

    @classmethod
    @redis_fallback(default_return=True)
    def set_allowed_users(
        cls,
        tenant_id: str,
        document_id: str,
        user_ids: list[str],
    ) -> bool:
        """Set list of users allowed to download the document.

        Args:
            tenant_id: Tenant ID
            document_id: Document ID
            user_ids: List of user IDs allowed to download (empty list means everyone can download)

        Returns:
            Whether the operation was successful
        """
        key = cls._get_permission_key(tenant_id, document_id)

        # Delete old configuration
        redis_client.delete(key)

        # If user_ids is empty, no restriction is applied
        if not user_ids:
            logger.info(f"Document {document_id} permission restriction removed")
            return True

        # Add new user ID set
        redis_client.sadd(key, *user_ids)

        # Set expiration time
        redis_client.expire(key, cls.RETENTION_SECONDS)

        logger.info(f"Document {document_id} permission set for {len(user_ids)} users")
        return True

    @classmethod
    @redis_fallback(default_return=set())
    def get_allowed_users(
        cls,
        tenant_id: str,
        document_id: str,
    ) -> set[str]:
        """Get set of user IDs allowed to download.

        Returns:
            Set of user IDs. Empty set means no restriction.
        """
        key = cls._get_permission_key(tenant_id, document_id)
        members = redis_client.smembers(key)

        # Decode byte strings
        return {m.decode("utf-8") if isinstance(m, bytes) else m for m in members}

    @classmethod
    @redis_fallback(default_return=(True, None))
    def check_download_permission(
        cls,
        tenant_id: str,
        document_id: str,
        user_id: str,
    ) -> tuple[bool, Optional[str]]:
        """Check if user has download permission.

        Returns:
            (has_permission, error_message)
        """
        allowed_users = cls.get_allowed_users(tenant_id, document_id)

        # If no permission configured (empty set), everyone can download
        if not allowed_users:
            return True, None

        # Check if user is in allowed list
        if user_id not in allowed_users:
            logger.warning(f"User {user_id} denied download access to document {document_id}")
            return False, "You don't have permission to download this document"

        return True, None

    # ========== Download Logging ==========

    @classmethod
    def _get_log_key(cls, tenant_id: str, document_id: str) -> str:
        """Get Redis key for download logs."""
        return f"{cls.LOG_PREFIX}:{tenant_id}:{document_id}"

    @classmethod
    def _get_count_key(cls, tenant_id: str, document_id: str, user_id: str) -> str:
        """Get Redis key for user download count."""
        return f"{cls.COUNT_PREFIX}:{tenant_id}:{document_id}:{user_id}"

    @classmethod
    @redis_fallback(default_return=False)
    def record_download(
        cls,
        tenant_id: str,
        document_id: str,
        user_id: str,
        user_name: str,
        ip: str,
    ) -> bool:
        """Record a document download.

        Args:
            tenant_id: Tenant ID
            document_id: Document ID
            user_id: User ID
            user_name: User name
            ip: User IP address

        Returns:
            Whether the operation was successful
        """
        log_key = cls._get_log_key(tenant_id, document_id)

        # Create download record
        record = {
            "user_id": user_id,
            "user_name": user_name,
            "download_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "ip": ip,
        }

        # Add to list head (newest first)
        redis_client.lpush(log_key, json.dumps(record, ensure_ascii=False))

        # Limit list length
        redis_client.ltrim(log_key, 0, cls.MAX_LOG_ENTRIES - 1)

        # Set expiration time
        redis_client.expire(log_key, cls.RETENTION_SECONDS)

        # Increment user download count
        count_key = cls._get_count_key(tenant_id, document_id, user_id)
        redis_client.incr(count_key)
        redis_client.expire(count_key, cls.RETENTION_SECONDS)

        logger.info(f"Recorded download: user={user_name}, document={document_id}, ip={ip}")
        return True

    @classmethod
    @redis_fallback(default_return=[])
    def get_download_logs(
        cls,
        tenant_id: str,
        document_id: str,
        limit: int = 100,
    ) -> list[dict[str, Any]]:
        """Get download logs for a document.

        Args:
            tenant_id: Tenant ID
            document_id: Document ID
            limit: Maximum number of records to return

        Returns:
            List of download records
        """
        log_key = cls._get_log_key(tenant_id, document_id)

        # Get latest N records
        records = redis_client.lrange(log_key, 0, limit - 1)

        result = []
        for record_json in records:
            try:
                if isinstance(record_json, bytes):
                    record_json = record_json.decode("utf-8")
                record_data = json.loads(record_json)
                result.append(record_data)
            except Exception as e:
                logger.warning(f"Failed to parse download log: {e}")
                continue

        return result

    @classmethod
    @redis_fallback(default_return=0)
    def get_user_download_count(
        cls,
        tenant_id: str,
        document_id: str,
        user_id: str,
    ) -> int:
        """Get user's download count for a specific document."""
        count_key = cls._get_count_key(tenant_id, document_id, user_id)
        count = redis_client.get(count_key)
        return int(count) if count else 0

    @classmethod
    @redis_fallback(default_return=False)
    def clear_permission(
        cls,
        tenant_id: str,
        document_id: str,
    ) -> bool:
        """Clear download permission configuration for a document."""
        key = cls._get_permission_key(tenant_id, document_id)
        redis_client.delete(key)
        logger.info(f"Cleared permission for document {document_id}")
        return True