pdf/gateway/app/services/store.py

import threading
import uuid
from datetime import UTC, datetime
from typing import Any


def extract_permissions(doc_instance: Any) -> dict[str, Any] | None:
    """Read the engine's DocumentPermissions off a doc into a plain dict, or None
    if unavailable (treated downstream as unrestricted)."""
    try:
        p = doc_instance.permissions
        return {
            "isEncrypted": p.is_encrypted,
            "encryption": p.encryption,
            "securityRevision": p.security_revision,
            "ownerUnlocked": p.owner_unlocked,
            "canPrint": p.can_print,
            "canPrintHighRes": p.can_print_high_res,
            "canModify": p.can_modify,
            "canCopy": p.can_copy,
            "canAnnotate": p.can_annotate,
            "canFillForms": p.can_fill_forms,
            "canExtractForAccessibility": p.can_extract_for_accessibility,
            "canAssemble": p.can_assemble,
        }
    except Exception:
        return None


class DocumentStore:
    def __init__(self):
        self._lock = threading.Lock()
        self._documents: dict[str, dict[str, Any]] = {}

    def add_document(
        self,
        filename: str,
        bytes_data: bytes,
        doc_instance: Any,
        permissions: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        doc_id = str(uuid.uuid4())
        uploaded_at = datetime.now(UTC).isoformat().replace("+00:00", "Z")
        page_width = 612.0
        page_height = 792.0
        if doc_instance.page_count > 0:
            try:
                page_0 = doc_instance.get_page(0)
                page_width = page_0.width
                page_height = page_0.height
            except Exception:
                pass

        # Permissions are computed from the freshly-loaded doc at upload; on derived
        # docs (post-edit, which are saved decrypted) the caller passes the original's
        # permissions forward so enforcement stays consistent.
        if permissions is None:
            permissions = extract_permissions(doc_instance)

        info = {
            "id": doc_id,
            "filename": filename,
            "sizeBytes": len(bytes_data),
            "totalPages": doc_instance.page_count,
            "pageWidth": page_width,
            "pageHeight": page_height,
            "uploadedAt": uploaded_at,
            "status": "ready",
            "doc_instance": doc_instance,
            "bytes_data": bytes_data,
            "permissions": permissions,
        }

        with self._lock:
            self._documents[doc_id] = info

        return info

    def get_document(self, doc_id: str) -> dict[str, Any] | None:
        with self._lock:
            return self._documents.get(doc_id)

    def list_documents(self) -> list[dict[str, Any]]:
        with self._lock:
            return list(self._documents.values())

    def delete_document(self, doc_id: str) -> bool:
        with self._lock:
            if doc_id in self._documents:
                del self._documents[doc_id]
                return True
            return False


document_store = DocumentStore()