肉末
83 lines
2.0 KiB
Python
83 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import HTTPException, UploadFile
|
|
|
|
from app import config
|
|
|
|
ALLOWED_EXTENSIONS = frozenset(
|
|
{
|
|
"png",
|
|
"jpg",
|
|
"jpeg",
|
|
"gif",
|
|
"webp",
|
|
"svg",
|
|
"txt",
|
|
"md",
|
|
"pdf",
|
|
"doc",
|
|
"docx",
|
|
"xls",
|
|
"xlsx",
|
|
"ppt",
|
|
"pptx",
|
|
"csv",
|
|
"zip",
|
|
"7z",
|
|
"tar",
|
|
"gz",
|
|
"rar",
|
|
"json",
|
|
}
|
|
)
|
|
|
|
_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")
|
|
|
|
|
|
def _safe_original_name(filename: str | None) -> str:
|
|
name = Path(filename or "file").name
|
|
name = _SAFE_NAME_RE.sub("_", name).strip("._") or "file"
|
|
return name[:255]
|
|
|
|
|
|
def _extension(filename: str) -> str:
|
|
return Path(filename).suffix.lstrip(".").lower()
|
|
|
|
|
|
async def save_upload(file: UploadFile) -> tuple[str, int, str, str]:
|
|
original = _safe_original_name(file.filename)
|
|
ext = _extension(original)
|
|
if ext not in ALLOWED_EXTENSIONS:
|
|
raise HTTPException(status_code=400, detail="File type not allowed")
|
|
|
|
max_bytes = config.settings.max_upload_mb * 1024 * 1024
|
|
config.settings.uploads_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
stored_name = f"{uuid.uuid4().hex}.{ext}"
|
|
dest = config.settings.uploads_dir / stored_name
|
|
size = 0
|
|
try:
|
|
with dest.open("wb") as out:
|
|
while True:
|
|
chunk = await file.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
size += len(chunk)
|
|
if size > max_bytes:
|
|
raise HTTPException(status_code=413, detail="File too large")
|
|
out.write(chunk)
|
|
except HTTPException:
|
|
dest.unlink(missing_ok=True)
|
|
raise
|
|
except Exception:
|
|
dest.unlink(missing_ok=True)
|
|
raise
|
|
|
|
mime = file.content_type or "application/octet-stream"
|
|
rel_path = f"uploads/{stored_name}"
|
|
return rel_path, size, mime, original
|