nocobase-mcp/nocobase_crm_mcp/uploads.py

from __future__ import annotations

import fnmatch
import tempfile
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any

from starlette.datastructures import UploadFile
from starlette.requests import Request
from starlette.responses import JSONResponse

from .nocobase_http import NocoBaseHttpError

_COPY_CHUNK_BYTES = 1024 * 1024
_MULTIPART_OVERHEAD_BYTES = 64 * 1024


class _UploadTooLarge(ValueError):
    pass


class UploadEndpoint:
    """ASGI endpoint that forwards one bounded multipart file to NocoBase."""

    def __init__(
        self,
        http_client: Any,
        max_bytes: int,
        temp_dir: str | Path,
        allowed_mime_types: Sequence[str] | None = None,
    ) -> None:
        self._http_client = http_client
        self._max_bytes = max_bytes
        self._temp_dir = Path(temp_dir)
        self._allowed_mime_types = tuple(allowed_mime_types) if allowed_mime_types is not None else None

    async def __call__(self, request: Request) -> JSONResponse:
        if request.method != "POST":
            return _error_response("invalid_upload", "Uploads require POST", 405)
        if not _is_multipart(request):
            return _error_response("invalid_upload", "Uploads require multipart form data", 400)

        try:
            request = await _bounded_request(request, self._max_bytes + _MULTIPART_OVERHEAD_BYTES)
            form = await request.form()
        except _UploadTooLarge:
            return _error_response("upload_too_large", "Upload exceeds the configured size limit", 413)
        except Exception:
            return _error_response("invalid_upload", "Could not read upload form data", 400)

        items = list(form.multi_items())
        uploads = [(name, value) for name, value in items if isinstance(value, UploadFile)]
        if len(items) != 1 or len(uploads) != 1 or uploads[0][0] != "file":
            await _close_uploads(value for _, value in uploads)
            return _error_response("invalid_upload", "Provide exactly one file upload", 400)

        upload = uploads[0][1]
        if not isinstance(upload, UploadFile):
            return _error_response("invalid_upload", "Provide exactly one file upload", 400)

        filename = upload.filename or ""
        mimetype = upload.content_type
        if self._allowed_mime_types is not None and not any(
            fnmatch.fnmatch(mimetype or "", allowed) for allowed in self._allowed_mime_types
        ):
            await upload.close()
            return _error_response("unsupported_media_type", "This media type is not allowed", 415)

        temporary_path: Path | None = None
        try:
            with tempfile.NamedTemporaryFile(delete=False, dir=self._temp_dir) as temporary:
                temporary_path = Path(temporary.name)
                size = 0
                while chunk := await upload.read(_COPY_CHUNK_BYTES):
                    size += len(chunk)
                    if size > self._max_bytes:
                        return _error_response("upload_too_large", "Upload exceeds the configured size limit", 413)
                    temporary.write(chunk)

            attachment = self._http_client.upload_attachment(temporary_path, filename, mimetype)
            attachment_id = _attachment_id(attachment)
            if attachment_id is None:
                return _error_response("nb_upload_failed", "NocoBase attachment upload failed", 502)
            return JSONResponse({"upload_id": attachment_id, "filename": filename, "size": size, "mimetype": mimetype})
        except NocoBaseHttpError as error:
            if error.authentication_failed:
                return _error_response("nb_upload_unauthorized", "NocoBase attachment upload failed", 502)
            if error.status_code == 403:
                return _error_response(
                    "nb_upload_forbidden",
                    "The current NocoBase role cannot upload project attachments",
                    403,
                )
            return _error_response("nb_upload_failed", "NocoBase attachment upload failed", 502)
        except Exception:
            return _error_response("nb_upload_failed", "NocoBase attachment upload failed", 502)
        finally:
            await upload.close()
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)


def _is_multipart(request: Request) -> bool:
    content_type = request.headers.get("content-type", "")
    return content_type.split(";", 1)[0].strip().lower() == "multipart/form-data"


async def _bounded_request(request: Request, max_body_bytes: int) -> Request:
    content_length = request.headers.get("content-length")
    if content_length is not None:
        try:
            if int(content_length) > max_body_bytes:
                raise _UploadTooLarge
        except ValueError:
            pass

    received_bytes = 0
    chunks: list[bytes] = []
    while True:
        message = await request.receive()
        if message.get("type") == "http.request":
            body = message.get("body", b"")
            received_bytes += len(body)
            if received_bytes > max_body_bytes:
                raise _UploadTooLarge
            chunks.append(body)
            if not message.get("more_body", False):
                break
        elif message.get("type") == "http.disconnect":
            break

    body = b"".join(chunks)
    replayed = False

    async def receive() -> dict[str, Any]:
        nonlocal replayed
        if replayed:
            return {"type": "http.disconnect"}
        replayed = True
        return {"type": "http.request", "body": body, "more_body": False}

    return Request(request.scope, receive=receive)


def _attachment_id(attachment: Any) -> int | None:
    if not isinstance(attachment, Mapping):
        return None
    attachment_id = attachment.get("id")
    if isinstance(attachment_id, int) and not isinstance(attachment_id, bool):
        return attachment_id
    return None


async def _close_uploads(uploads: Iterable[UploadFile]) -> None:
    for upload in uploads:
        await upload.close()


def _error_response(code: str, message: str, status_code: int) -> JSONResponse:
    return JSONResponse({"error": {"code": code, "message": message}}, status_code=status_code)