import hashlib
import re
from pathlib import Path
from fastapi import UploadFile, HTTPException, status
from .config import settings
# Magic-byte sniffing for the formats we accept; anything else falls back to
# the client-provided mime but still gets stored (documents may be arbitrary).
_SNIFFERS: list[tuple[bytes, str]] = [
(b"\x89PNG\r\n\x1a\n", "image/png"),
(b"\xff\xd8\xff", "image/jpeg"),
(b"GIF8", "image/gif"),
(b"\x1a\x45\xdf\xa3", "video/webm"), # EBML header (Matroska/WebM)
(b"%PDF-", "application/pdf"),
(b"RIFF", "image/webp"), # RIFF....WEBP — verified below
]
_SAFE_NAME = re.compile(r"[^A-Za-z0-9._\- ]+")
def sniff_mime(head: bytes, fallback: str) -> str:
for magic, mime in _SNIFFERS:
if head.startswith(magic):
if magic == b"RIFF" and head[8:12] != b"WEBP":
continue
return mime
return fallback or "application/octet-stream"
def safe_filename(name: str) -> str:
name = _SAFE_NAME.sub("_", name).strip() or "file"
return name[:120]
async def store_upload(upload: UploadFile) -> tuple[str, str, int, str]:
"""Persist an upload. Returns (sha256_hex, rel_path, size, mime).
Raises HTTPException 413 if the file exceeds the configured limit.
"""
sha = hashlib.sha256()
size = 0
head = b""
chunks: list[bytes] = []
while chunk := await upload.read(1024 * 1024):
if not head:
head = chunk[:32]
size += len(chunk)
if size > settings.max_upload_bytes:
raise HTTPException(status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, "File too large")
sha.update(chunk)
chunks.append(chunk)
if size == 0:
raise HTTPException(status.HTTP_400_BAD_REQUEST, "Empty file")
mime = sniff_mime(head, upload.content_type or "")
digest = sha.hexdigest()
rel_path = f"{digest[:2]}/{digest[2:4]}/{digest}"
dest = settings.files_dir / rel_path
if not dest.exists(): # content-addressed dedup
dest.parent.mkdir(parents=True, exist_ok=True)
dest.write_bytes(b"".join(chunks))
return digest, rel_path, size, mime