189 lines
7.5 KiB
Python
189 lines
7.5 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
from fastapi import HTTPException, UploadFile
|
|
|
|
from app.config import settings
|
|
from app.enums.admin_upload import (
|
|
ADMIN_UPLOAD_AUDIO_EXTENSIONS,
|
|
ADMIN_UPLOAD_AUDIO_MAX_BYTES,
|
|
ADMIN_UPLOAD_FILE_MAX_BYTES,
|
|
ADMIN_UPLOAD_FILE_NAME_MAX_LEN,
|
|
ADMIN_UPLOAD_IMAGE_EXTENSIONS,
|
|
ADMIN_UPLOAD_IMAGE_MAX_BYTES,
|
|
ADMIN_UPLOAD_PDF_EXTENSIONS,
|
|
ADMIN_UPLOAD_PDF_MAX_BYTES,
|
|
ADMIN_UPLOAD_VIDEO_EXTENSIONS,
|
|
ADMIN_UPLOAD_VIDEO_MAX_BYTES,
|
|
AdminUploadResourceTypeEnum,
|
|
AdminUploadSceneEnum,
|
|
)
|
|
|
|
CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class StoredAdminUploadFile:
|
|
storage_path: str
|
|
file_url: str
|
|
file_name: str
|
|
original_file_name: str | None
|
|
file_size_bytes: int
|
|
|
|
|
|
def _upload_root() -> Path:
|
|
return Path(settings.UPLOAD_LOCAL_PATH).resolve()
|
|
|
|
|
|
def _safe_original_name(filename: str | None) -> str | None:
|
|
if not filename:
|
|
return None
|
|
name = Path(filename).name.strip()
|
|
name = re.sub(r"[\\/\r\n\t]+", "_", name)
|
|
return name[:ADMIN_UPLOAD_FILE_NAME_MAX_LEN] or None
|
|
|
|
|
|
def _ext(filename: str | None) -> str:
|
|
return Path(filename or "").suffix.lower()
|
|
|
|
|
|
def _allowed_extensions(resource_type: AdminUploadResourceTypeEnum) -> set[str] | None:
|
|
if resource_type == AdminUploadResourceTypeEnum.IMAGE:
|
|
return ADMIN_UPLOAD_IMAGE_EXTENSIONS
|
|
if resource_type == AdminUploadResourceTypeEnum.VIDEO:
|
|
return ADMIN_UPLOAD_VIDEO_EXTENSIONS
|
|
if resource_type == AdminUploadResourceTypeEnum.AUDIO:
|
|
return ADMIN_UPLOAD_AUDIO_EXTENSIONS
|
|
if resource_type == AdminUploadResourceTypeEnum.PDF:
|
|
return ADMIN_UPLOAD_PDF_EXTENSIONS
|
|
return None
|
|
|
|
|
|
def _max_bytes(resource_type: AdminUploadResourceTypeEnum) -> int:
|
|
if resource_type == AdminUploadResourceTypeEnum.IMAGE:
|
|
return ADMIN_UPLOAD_IMAGE_MAX_BYTES
|
|
if resource_type == AdminUploadResourceTypeEnum.VIDEO:
|
|
return ADMIN_UPLOAD_VIDEO_MAX_BYTES
|
|
if resource_type == AdminUploadResourceTypeEnum.AUDIO:
|
|
return ADMIN_UPLOAD_AUDIO_MAX_BYTES
|
|
if resource_type == AdminUploadResourceTypeEnum.PDF:
|
|
return ADMIN_UPLOAD_PDF_MAX_BYTES
|
|
return ADMIN_UPLOAD_FILE_MAX_BYTES
|
|
|
|
|
|
def validate_file_name(filename: str | None, resource_type: AdminUploadResourceTypeEnum) -> str:
|
|
safe_name = _safe_original_name(filename)
|
|
if not safe_name:
|
|
raise HTTPException(status_code=400, detail="请选择文件")
|
|
allowed = _allowed_extensions(resource_type)
|
|
ext = _ext(safe_name)
|
|
if allowed is not None and ext not in allowed:
|
|
if resource_type == AdminUploadResourceTypeEnum.IMAGE:
|
|
raise HTTPException(status_code=400, detail="仅支持 jpg/jpeg/png/webp/gif 图片")
|
|
if resource_type == AdminUploadResourceTypeEnum.VIDEO:
|
|
raise HTTPException(status_code=400, detail="仅支持 mp4/mov/m4v/webm 视频")
|
|
if resource_type == AdminUploadResourceTypeEnum.AUDIO:
|
|
raise HTTPException(status_code=400, detail="仅支持 mp3/wav/m4a/aac 音频")
|
|
if resource_type == AdminUploadResourceTypeEnum.PDF:
|
|
raise HTTPException(status_code=400, detail="仅支持 PDF 文件")
|
|
return safe_name
|
|
|
|
|
|
def _kind_dir(resource_type: AdminUploadResourceTypeEnum) -> str:
|
|
if resource_type == AdminUploadResourceTypeEnum.IMAGE:
|
|
return "images"
|
|
if resource_type == AdminUploadResourceTypeEnum.VIDEO:
|
|
return "videos"
|
|
if resource_type == AdminUploadResourceTypeEnum.AUDIO:
|
|
return "audios"
|
|
return "files"
|
|
|
|
|
|
def _scene_base_dir(scene: AdminUploadSceneEnum, resource_type: AdminUploadResourceTypeEnum) -> Path:
|
|
if scene == AdminUploadSceneEnum.HOME_MATERIAL_REFERENCE:
|
|
if resource_type not in {AdminUploadResourceTypeEnum.IMAGE, AdminUploadResourceTypeEnum.VIDEO, AdminUploadResourceTypeEnum.AUDIO}:
|
|
raise HTTPException(status_code=400, detail="首页素材附件仅支持图片、视频、音频")
|
|
return Path("home_materials") / "references" / _kind_dir(resource_type)
|
|
if scene == AdminUploadSceneEnum.SYSTEM_LOGO:
|
|
if resource_type != AdminUploadResourceTypeEnum.IMAGE:
|
|
raise HTTPException(status_code=400, detail="系统Logo仅支持图片")
|
|
return Path("admin_uploads") / "system_logo" / "images"
|
|
if scene == AdminUploadSceneEnum.SYSTEM_PDF:
|
|
if resource_type != AdminUploadResourceTypeEnum.PDF:
|
|
raise HTTPException(status_code=400, detail="系统PDF仅支持PDF文件")
|
|
return Path("admin_uploads") / "system_pdf" / "files"
|
|
if scene == AdminUploadSceneEnum.OPEN_TYPE_THUMB:
|
|
if resource_type != AdminUploadResourceTypeEnum.IMAGE:
|
|
raise HTTPException(status_code=400, detail="开户方式缩略图仅支持图片")
|
|
return Path("admin_uploads") / "open_type_thumb" / "images"
|
|
return Path("admin_uploads") / "common" / _kind_dir(resource_type)
|
|
|
|
|
|
def build_target(scene: AdminUploadSceneEnum, resource_type: AdminUploadResourceTypeEnum, admin_id: str, original_filename: str | None) -> tuple[Path, str, str]:
|
|
safe_original = validate_file_name(original_filename, resource_type)
|
|
now = datetime.now()
|
|
date_dir = now.strftime("%Y/%m/%d")
|
|
timestamp = now.strftime("%Y%m%d_%H%M%S")
|
|
suffix = uuid.uuid4().hex[:8]
|
|
ext = _ext(safe_original)
|
|
prefix = {
|
|
AdminUploadResourceTypeEnum.IMAGE: "admin_img",
|
|
AdminUploadResourceTypeEnum.VIDEO: "admin_video",
|
|
AdminUploadResourceTypeEnum.AUDIO: "admin_audio",
|
|
AdminUploadResourceTypeEnum.PDF: "admin_pdf",
|
|
}.get(resource_type, "admin_file")
|
|
file_name = f"{prefix}_{admin_id}_{timestamp}_{suffix}{ext}"
|
|
rel_dir = _scene_base_dir(scene, resource_type) / date_dir
|
|
storage_path = _upload_root() / rel_dir / file_name
|
|
file_url = f"/uploads/{(rel_dir / file_name).as_posix()}"
|
|
return storage_path, file_url, file_name
|
|
|
|
|
|
async def save_upload_file(file: UploadFile, *, scene: AdminUploadSceneEnum, resource_type: AdminUploadResourceTypeEnum, admin_id: str) -> StoredAdminUploadFile:
|
|
original_file_name = validate_file_name(file.filename, resource_type)
|
|
max_bytes = _max_bytes(resource_type)
|
|
fd, temp_path = tempfile.mkstemp(prefix="admin_upload_", suffix=".tmp")
|
|
total = 0
|
|
final_path: Path | None = None
|
|
try:
|
|
with os.fdopen(fd, "wb") as out:
|
|
while True:
|
|
chunk = await file.read(CHUNK_SIZE)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > max_bytes:
|
|
raise HTTPException(status_code=400, detail=f"文件大小不能超过 {max_bytes // 1024 // 1024}MB")
|
|
out.write(chunk)
|
|
final_path, file_url, file_name = build_target(scene, resource_type, admin_id, original_file_name)
|
|
final_path.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.move(temp_path, final_path)
|
|
temp_path = ""
|
|
return StoredAdminUploadFile(
|
|
storage_path=str(final_path),
|
|
file_url=file_url,
|
|
file_name=file_name,
|
|
original_file_name=original_file_name,
|
|
file_size_bytes=total,
|
|
)
|
|
except Exception:
|
|
if temp_path:
|
|
try:
|
|
os.remove(temp_path)
|
|
except OSError:
|
|
pass
|
|
if final_path and final_path.exists():
|
|
try:
|
|
final_path.unlink()
|
|
except OSError:
|
|
pass
|
|
raise
|