100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
"""对象存储（S3 兼容）：上传/下载图片，含校验与去重。
|
||
|
||
boto3 是同步库，直接在 async 路径里调用会阻塞事件循环（传大图时整个服务卡住），
|
||
所以对外暴露 async 版本（to_thread 包装），同步版本保留给测试和脚本用。
|
||
"""
|
||
import asyncio
|
||
import hashlib
|
||
import io
|
||
import uuid
|
||
|
||
import boto3
|
||
from botocore.config import Config
|
||
from PIL import Image as PILImage
|
||
|
||
from app.config import settings
|
||
|
||
# Pillow 格式 → 扩展名
|
||
_FORMAT_EXT = {"JPEG": "jpg", "PNG": "png", "WEBP": "webp"}
|
||
# 解压炸弹防护：单边像素上限
|
||
_MAX_DIMENSION = 12000
|
||
|
||
|
||
class StorageError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class ImageValidationError(ValueError):
|
||
pass
|
||
|
||
|
||
def _s3_client():
|
||
if not settings.s3_bucket:
|
||
raise StorageError("S3 未配置（S3_BUCKET 为空）")
|
||
return boto3.client(
|
||
"s3",
|
||
endpoint_url=settings.s3_endpoint_url,
|
||
region_name=settings.s3_region,
|
||
aws_access_key_id=settings.s3_access_key,
|
||
aws_secret_access_key=settings.s3_secret_key,
|
||
config=Config(signature_version="s3v4"),
|
||
)
|
||
|
||
|
||
def validate_image(data: bytes) -> tuple[str, str]:
|
||
"""校验是真实图片且格式在白名单内。返回 (mime, ext)。"""
|
||
if len(data) > settings.max_upload_bytes:
|
||
raise ImageValidationError("文件超过大小上限")
|
||
try:
|
||
img = PILImage.open(io.BytesIO(data))
|
||
img.verify() # 校验完整性（魔数级）
|
||
# verify 后需重新打开才能读属性
|
||
img = PILImage.open(io.BytesIO(data))
|
||
except Exception as e: # noqa: BLE001
|
||
raise ImageValidationError(f"不是有效的图片文件：{e}") from e
|
||
|
||
fmt = img.format or ""
|
||
if fmt not in _FORMAT_EXT:
|
||
raise ImageValidationError(f"不支持的图片格式：{fmt}")
|
||
if img.width > _MAX_DIMENSION or img.height > _MAX_DIMENSION:
|
||
raise ImageValidationError("图片尺寸过大")
|
||
|
||
mime = f"image/{'jpeg' if fmt == 'JPEG' else fmt.lower()}"
|
||
if mime not in settings.allowed_image_mimes:
|
||
raise ImageValidationError(f"不支持的 MIME：{mime}")
|
||
return mime, _FORMAT_EXT[fmt]
|
||
|
||
|
||
def upload_image(user_id: str, data: bytes, ext: str, mime: str) -> str:
|
||
"""上传到 S3，返回 object_key。"""
|
||
key = f"{user_id}/{uuid.uuid4().hex}.{ext}"
|
||
client = _s3_client()
|
||
client.put_object(
|
||
Bucket=settings.s3_bucket,
|
||
Key=key,
|
||
Body=data,
|
||
ContentType=mime,
|
||
)
|
||
return key
|
||
|
||
|
||
def download_image(object_key: str) -> bytes:
|
||
client = _s3_client()
|
||
resp = client.get_object(Bucket=settings.s3_bucket, Key=object_key)
|
||
return resp["Body"].read()
|
||
|
||
|
||
def sha256_hex(data: bytes) -> str:
|
||
return hashlib.sha256(data).hexdigest()
|
||
|
||
|
||
# ---- async 包装：在 async 路径里请用这两个 ----
|
||
|
||
|
||
async def upload_image_async(user_id: str, data: bytes, ext: str, mime: str) -> str:
|
||
return await asyncio.to_thread(upload_image, user_id, data, ext, mime)
|
||
|
||
|
||
async def download_image_async(object_key: str) -> bytes:
|
||
return await asyncio.to_thread(download_image, object_key)
|