113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
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
|
|
|
|
|
|
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:
|
|
form = await request.form()
|
|
except Exception:
|
|
return _error_response("invalid_upload", "Could not read upload form data", 400)
|
|
|
|
uploads = [(name, value) for name, value in form.multi_items() if isinstance(value, UploadFile)]
|
|
file_values = form.getlist("file")
|
|
if len(uploads) != 1 or len(file_values) != 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 = file_values[0]
|
|
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:
|
|
code = "nb_upload_unauthorized" if error.status_code in {401, 403} else "nb_upload_failed"
|
|
return _error_response(code, "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"
|
|
|
|
|
|
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)
|