Newer
Older
bugtrail / server / app / storage.py
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"%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