app.py

import os
import pathlib
import uuid

from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse

DATA_DIR = pathlib.Path(os.environ.get("DATA_DIR", "/tmp/uploads"))
DATA_DIR.mkdir(parents=True, exist_ok=True)
MAX_BYTES = int(os.environ.get("MAX_BYTES", str(30 * 1024 * 1024)))
# Optional shared upload token. If UPLOAD_TOKEN is unset, uploads are open (temp use).
UPLOAD_TOKEN = os.environ.get("UPLOAD_TOKEN", "").strip()
# Optional public base override, e.g. https://looknicemm1-tmp-files.hf.space
PUBLIC_BASE = os.environ.get("PUBLIC_BASE", "").strip().rstrip("/")

EXT_BY_TYPE = {
    "image/jpeg": "jpg",
    "image/jpg": "jpg",
    "image/png": "png",
    "image/webp": "webp",
    "image/gif": "gif",
    "video/mp4": "mp4",
    "video/quicktime": "mov",
    "video/webm": "webm",
    "audio/mpeg": "mp3",
    "audio/mp4": "m4a",
    "audio/wav": "wav",
    "audio/x-wav": "wav",
    "audio/aac": "aac",
    "audio/ogg": "ogg",
}

app = FastAPI(title="tmp-file-service")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def public_base(request: Request) -> str:
    if PUBLIC_BASE:
        return PUBLIC_BASE
    host = request.headers.get("x-forwarded-host") or request.headers.get("host") or "localhost"
    proto = request.headers.get("x-forwarded-proto", "https")
    return f"{proto}://{host}"


def check_auth(request: Request) -> None:
    if not UPLOAD_TOKEN:
        return
    header = request.headers.get("authorization", "")
    token = header[7:].strip() if header.lower().startswith("bearer ") else request.query_params.get("token", "")
    if token != UPLOAD_TOKEN:
        raise HTTPException(401, "invalid upload token")


def pick_ext(file: UploadFile) -> str:
    ext = EXT_BY_TYPE.get((file.content_type or "").lower())
    if ext:
        return ext
    suffix = pathlib.Path(file.filename or "").suffix.lstrip(".").lower()
    return suffix or "bin"


@app.get("/")
def health():
    return {"ok": True, "service": "tmp-file-service", "auth": bool(UPLOAD_TOKEN)}


@app.post("/upload")
async def upload(request: Request, file: UploadFile = File(...)):
    check_auth(request)
    data = await file.read()
    if not data:
        raise HTTPException(400, "empty file")
    if len(data) > MAX_BYTES:
        raise HTTPException(413, f"file too large (> {MAX_BYTES} bytes)")
    name = f"{uuid.uuid4().hex}.{pick_ext(file)}"
    (DATA_DIR / name).write_bytes(data)
    url = f"{public_base(request)}/files/{name}"
    return {"success": True, "url": url, "filename": name, "size": len(data)}


@app.get("/files/{name}")
def get_file(name: str):
    if "/" in name or ".." in name:
        raise HTTPException(400, "bad name")
    path = DATA_DIR / name
    if not path.is_file():
        raise HTTPException(404, "not found")
    return FileResponse(path)