101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
"""Authenticated binary assets used by annotations."""
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
|
|
from fastapi.responses import FileResponse
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.auth.dependencies import get_current_user, verify_csrf
|
|
from app.config import settings
|
|
from app.db import get_db
|
|
from app.models.document import Document
|
|
|
|
router = APIRouter(prefix="/documents", tags=["Assets"])
|
|
_ASSET_CHUNK_SIZE = 1024 * 1024
|
|
_MAX_ASSET_BYTES = 25 * 1024 * 1024
|
|
|
|
|
|
def get_asset_path(document_id: str, ref: str, *, create: bool = False) -> Path:
|
|
"""Return a safe asset path; only writes are allowed to create directories."""
|
|
try:
|
|
asset_ref = str(uuid.UUID(ref))
|
|
except ValueError as err:
|
|
raise HTTPException(status_code=404, detail="Asset not found") from err
|
|
|
|
directory = settings.PDF_STORAGE_PATH / "assets" / document_id
|
|
if create:
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return directory / asset_ref
|
|
|
|
|
|
def _is_supported_image(filepath: Path) -> bool:
|
|
with filepath.open("rb") as file:
|
|
header = file.read(12)
|
|
return (
|
|
header.startswith(b"\x89PNG\r\n\x1a\n")
|
|
or header.startswith(b"\xff\xd8\xff")
|
|
or header.startswith((b"GIF87a", b"GIF89a"))
|
|
or (header.startswith(b"RIFF") and header[8:12] == b"WEBP")
|
|
)
|
|
|
|
|
|
@router.post("/{id}/assets", dependencies=[Depends(verify_csrf)])
|
|
async def upload_asset(
|
|
id: str,
|
|
file: UploadFile = File(...),
|
|
user_id: int = Depends(get_current_user),
|
|
db: Session = Depends(get_db),
|
|
) -> dict[str, str]:
|
|
"""Upload a PNG, JPEG, GIF, or WebP annotation asset."""
|
|
document = db.scalar(select(Document).where(Document.id == id))
|
|
if not document:
|
|
raise HTTPException(status_code=404, detail="Document not found")
|
|
|
|
ref = str(uuid.uuid4())
|
|
filepath = get_asset_path(id, ref, create=True)
|
|
total_bytes = 0
|
|
|
|
try:
|
|
with filepath.open("wb") as destination:
|
|
while chunk := file.file.read(_ASSET_CHUNK_SIZE):
|
|
total_bytes += len(chunk)
|
|
if total_bytes > _MAX_ASSET_BYTES:
|
|
raise HTTPException(status_code=413, detail="Image asset is too large")
|
|
destination.write(chunk)
|
|
except HTTPException:
|
|
filepath.unlink(missing_ok=True)
|
|
raise
|
|
except Exception as err:
|
|
filepath.unlink(missing_ok=True)
|
|
raise HTTPException(status_code=500, detail="Failed to save asset") from err
|
|
|
|
if not _is_supported_image(filepath):
|
|
filepath.unlink(missing_ok=True)
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Only PNG, JPEG, GIF, and WebP images are supported",
|
|
)
|
|
|
|
return {"ref": ref, "url": f"/api/v1/documents/{id}/assets/{ref}"}
|
|
|
|
|
|
@router.get("/{id}/assets/{ref}")
|
|
async def get_asset(
|
|
id: str,
|
|
ref: str,
|
|
db: Session = Depends(get_db),
|
|
user_id: int = Depends(get_current_user),
|
|
) -> FileResponse:
|
|
"""Serve an annotation asset to an authenticated session."""
|
|
document = db.scalar(select(Document).where(Document.id == id))
|
|
if not document:
|
|
raise HTTPException(status_code=404, detail="Document not found")
|
|
|
|
filepath = get_asset_path(id, ref)
|
|
if not filepath.exists() or not filepath.is_file():
|
|
raise HTTPException(status_code=404, detail="Asset not found")
|
|
|
|
return FileResponse(filepath)
|