document-management/backend/dms/services/file_read_service.py

"""统一文档预览、下载授权及下载计数审计。"""

from __future__ import annotations

import logging
from pathlib import Path

from flask import current_app, send_file
from sqlalchemy import select

from dms.common.enums import AuditAction, AuditTarget, DocumentType
from dms.common.errors import (
    DocumentDownloadForbiddenError,
    FileNotFoundError as DmsFileNotFoundError,
    InternalError,
    PreviewUnavailableError,
    ResourceNotFoundError,
    SecurityLevelForbiddenError,
)
from dms.database.transaction import transaction
from dms.extensions import db
from dms.models import Document
from dms.security.auth_context import get_auth_context
from dms.services.audit_service import business_audit
from dms.services.authorization_service import evaluate_plan_access
from dms.services.document_query_service import _require_plan_access
from dms.storage.paths import UnsafeStoragePathError, resolve_storage_path

logger = logging.getLogger(__name__)


def _document(document_id: int) -> Document:
    document = db.session.scalar(
        select(Document).where(
            Document.id == document_id,
            Document.is_deleted.is_(False),
            Document.document_type.in_(
                [
                    DocumentType.MAIN.value,
                    DocumentType.SUB_PLAN.value,
                    DocumentType.ATTACHMENT.value,
                ]
            ),
        )
    )
    if document is None:
        raise ResourceNotFoundError("文档不存在")
    return document


def _authorize_view(document: Document) -> None:
    if document.document_type != DocumentType.ATTACHMENT.value:
        _require_plan_access(document, get_auth_context())


def _authorize_download(document: Document) -> None:
    if document.document_type == DocumentType.ATTACHMENT.value:
        return
    access = evaluate_plan_access(document, get_auth_context())
    if access.reason == "SECURITY":
        raise SecurityLevelForbiddenError()
    if not access.allowed or not access.can_download:
        raise DocumentDownloadForbiddenError()


def _path(document: Document) -> Path:
    try:
        path = resolve_storage_path(
            document.file_relative_path,
            current_app.config["DMS_STORAGE_ROOT"],
        )
    except UnsafeStoragePathError as exc:
        logger.error("检测到越界文档存储路径：document_id=%s", document.id)
        raise InternalError("文档存储路径异常") from exc
    if not path.is_file():
        raise DmsFileNotFoundError()
    try:
        with path.open("rb") as stream:
            stream.read(1)
    except OSError as exc:
        logger.exception("文档文件读取失败：document_id=%s", document.id)
        raise InternalError("文档文件读取失败") from exc
    return path


def _download_name(value: str) -> str:
    safe = value.replace("\r", "").replace("\n", "").replace("\x00", "")
    return Path(safe).name or "download"


def preview_document(document_id: int):
    document = _document(document_id)
    _authorize_view(document)
    if document.file_extension.lower() != "pdf":
        raise PreviewUnavailableError()
    path = _path(document)
    response = send_file(
        path,
        mimetype="application/pdf",
        as_attachment=False,
        download_name=_download_name(document.original_file_name),
        conditional=True,
    )
    response.headers["X-Content-Type-Options"] = "nosniff"
    return response


def _record_download(document: Document) -> None:
    document_id = document.id
    document_name = document.document_name
    target = (
        AuditTarget.ATTACHMENT
        if document.document_type == DocumentType.ATTACHMENT.value
        else AuditTarget.DOCUMENT
    )
    db.session.rollback()
    try:
        with transaction() as session:
            current = session.scalar(
                select(Document)
                .where(
                    Document.id == document_id,
                    Document.is_deleted.is_(False),
                )
                .with_for_update()
            )
            if current is None:
                raise ResourceNotFoundError("文档不存在")
            current.download_count += 1
            session.add(
                business_audit(
                    action=AuditAction.DOWNLOAD_DOCUMENT,
                    target=target,
                    target_id=document_id,
                    target_name=document_name,
                    detail={"documentType": current.document_type},
                )
            )
    except Exception:
        db.session.rollback()
        logger.exception("文档下载计数或审计写入失败：document_id=%s", document_id)


def download_document(document_id: int):
    document = _document(document_id)
    _authorize_download(document)
    path = _path(document)
    mime_type = document.mime_type or "application/octet-stream"
    download_name = _download_name(document.original_file_name)
    _record_download(document)
    response = send_file(
        path,
        mimetype=mime_type,
        as_attachment=True,
        download_name=download_name,
        conditional=True,
    )
    response.headers["X-Content-Type-Options"] = "nosniff"
    return response