237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
"""图片/文件存储抽象：本地文件系统（开发兜底）+ 七牛（生产）。"""
|
||
from __future__ import annotations
|
||
|
||
import mimetypes
|
||
import shutil
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
|
||
from config import get_settings
|
||
|
||
# 本地存储根目录（仓库根 data/media/）
|
||
_LOCAL_ROOT = Path(__file__).resolve().parents[2] / "data" / "media"
|
||
|
||
|
||
class ResourceSizeLimitError(ValueError):
|
||
"""上传、下载或本地读取超过单文件上限。"""
|
||
|
||
def __init__(self, limit_bytes: int, actual_bytes: int | None = None) -> None:
|
||
self.limit_bytes = limit_bytes
|
||
self.actual_bytes = actual_bytes
|
||
detail = f"图片文件超过大小上限（最多 {limit_bytes // 1024 // 1024} MB）"
|
||
if actual_bytes is not None:
|
||
detail += f"，当前约 {actual_bytes / 1024 / 1024:.1f} MB"
|
||
super().__init__(detail)
|
||
|
||
|
||
def _ext_from_content_type(content_type: str) -> str:
|
||
ctype = (content_type or "").split(";")[0].strip().lower()
|
||
mapping = {
|
||
"image/jpeg": ".jpg",
|
||
"image/png": ".png",
|
||
"image/webp": ".webp",
|
||
"image/gif": ".gif",
|
||
"image/bmp": ".bmp",
|
||
"image/heic": ".heic",
|
||
"video/mp4": ".mp4",
|
||
}
|
||
return mapping.get(ctype, ".jpg")
|
||
|
||
|
||
def image_extension(data: bytes, content_type: str = "") -> str:
|
||
"""按真实图片字节优先推断扩展名，避免 provider 错报 Content-Type。"""
|
||
if data[:8] == b"\x89PNG\r\n\x1a\n":
|
||
return ".png"
|
||
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||
return ".webp"
|
||
if data[:3] == b"GIF":
|
||
return ".gif"
|
||
return _ext_from_content_type(content_type)
|
||
|
||
|
||
async def download_bytes(
|
||
url: str,
|
||
referer: str | None = None,
|
||
timeout: float = 60.0,
|
||
max_bytes: int | None = None,
|
||
) -> tuple[bytes, str]:
|
||
"""流式下载远程图片并限制体积。返回 (bytes, content_type)。"""
|
||
limit = max_bytes or get_settings().image_download_max_bytes
|
||
if limit <= 0:
|
||
raise ValueError("图片下载大小上限必须大于 0")
|
||
headers = {"Referer": referer} if referer else {}
|
||
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
|
||
async with client.stream("GET", url, headers=headers) as resp:
|
||
resp.raise_for_status()
|
||
length = resp.headers.get("content-length")
|
||
if length:
|
||
try:
|
||
declared_size = int(length)
|
||
except ValueError:
|
||
declared_size = None
|
||
if declared_size is not None and declared_size > limit:
|
||
raise ResourceSizeLimitError(limit, declared_size)
|
||
|
||
chunks: list[bytes] = []
|
||
total = 0
|
||
async for chunk in resp.aiter_bytes():
|
||
total += len(chunk)
|
||
if total > limit:
|
||
raise ResourceSizeLimitError(limit, total)
|
||
chunks.append(chunk)
|
||
ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
|
||
return b"".join(chunks), ctype
|
||
|
||
|
||
def read_local_bytes(path: Path, max_bytes: int | None = None) -> bytes:
|
||
"""受单文件上限保护地读取本地媒体文件。"""
|
||
limit = max_bytes or get_settings().image_download_max_bytes
|
||
size = path.stat().st_size
|
||
if size > limit:
|
||
raise ResourceSizeLimitError(limit, size)
|
||
return path.read_bytes()
|
||
|
||
|
||
class LocalStorage:
|
||
"""开发期：落 data/media/，由 FastAPI /media 静态托管，返回 app_base_url 可访问 URL。"""
|
||
|
||
async def save_from_url(self, url: str, key_prefix: str = "", referer: str | None = None) -> str:
|
||
data, ctype = await download_bytes(url, referer)
|
||
key = f"{key_prefix.rstrip('/') + '/' if key_prefix else ''}{uuid.uuid4().hex}{image_extension(data, ctype)}"
|
||
return await self.save_bytes(data, key, ctype)
|
||
|
||
async def save_bytes(self, data: bytes, key: str, content_type: str = "") -> str:
|
||
return self.public_url(self._write(data, key, content_type))
|
||
|
||
def _write(self, data: bytes, key: str, content_type: str) -> str:
|
||
_LOCAL_ROOT.mkdir(parents=True, exist_ok=True)
|
||
safe_key = key.strip().lstrip("/")
|
||
if not safe_key:
|
||
safe_key = uuid.uuid4().hex
|
||
if not Path(safe_key).suffix:
|
||
safe_key += image_extension(data, content_type)
|
||
path = (_LOCAL_ROOT / safe_key).resolve()
|
||
if _LOCAL_ROOT.resolve() not in path.parents:
|
||
raise ValueError("非法存储 key")
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_bytes(data)
|
||
return safe_key
|
||
|
||
def public_url(self, key: str) -> str:
|
||
settings = get_settings()
|
||
return f"{settings.app_base_url.rstrip('/')}/media/{key}"
|
||
|
||
async def delete(self, key: str) -> None:
|
||
path = (_LOCAL_ROOT / key.strip().lstrip("/")).resolve()
|
||
if _LOCAL_ROOT.resolve() not in path.parents:
|
||
raise ValueError("非法存储 key")
|
||
if path.is_file():
|
||
path.unlink()
|
||
|
||
async def delete_prefix(self, prefix: str) -> int:
|
||
root = (_LOCAL_ROOT / prefix.strip().strip("/")).resolve()
|
||
if _LOCAL_ROOT.resolve() not in root.parents and root != _LOCAL_ROOT.resolve():
|
||
raise ValueError("非法存储前缀")
|
||
if not root.exists():
|
||
return 0
|
||
if root.is_file():
|
||
root.unlink()
|
||
return 1
|
||
count = sum(1 for path in root.rglob("*") if path.is_file())
|
||
shutil.rmtree(root)
|
||
return count
|
||
|
||
|
||
class QiniuStorage:
|
||
"""生产：上传七牛，返回绑定域名公网 URL（Ozon 可拉取）。"""
|
||
|
||
def _client(self):
|
||
import qiniu
|
||
|
||
settings = get_settings()
|
||
return qiniu.Auth(settings.qiniu_access_key, settings.qiniu_secret_key), settings
|
||
|
||
async def save_from_url(self, url: str, key_prefix: str = "", referer: str | None = None) -> str:
|
||
data, ctype = await download_bytes(url, referer)
|
||
return await self.save_bytes(data, f"{key_prefix}/{uuid.uuid4().hex}{_ext_from_content_type(ctype)}", ctype)
|
||
|
||
async def save_bytes(self, data: bytes, key: str, content_type: str = "") -> str:
|
||
import qiniu
|
||
|
||
auth, settings = self._client()
|
||
bucket = settings.qiniu_bucket
|
||
token = auth.upload_token(bucket, key, 3600)
|
||
ret, info = qiniu.put_data(token, key, data)
|
||
if info.status_code not in (200,):
|
||
raise RuntimeError(f"七牛上传失败：{info.error or info.text_body or info.status_code}")
|
||
return self.public_url(key)
|
||
|
||
def public_url(self, key: str) -> str:
|
||
settings = get_settings()
|
||
return f"{settings.qiniu_domain.rstrip('/')}/{key}"
|
||
|
||
async def delete(self, key: str) -> None:
|
||
import qiniu
|
||
|
||
auth, settings = self._client()
|
||
_, info = qiniu.BucketManager(auth).delete(settings.qiniu_bucket, key)
|
||
if info.status_code not in (200, 612): # 612：对象已不存在，视为删除完成
|
||
raise RuntimeError(f"七牛删除失败：{info.error or info.text_body or info.status_code}")
|
||
|
||
async def delete_prefix(self, prefix: str) -> int:
|
||
import qiniu
|
||
|
||
auth, settings = self._client()
|
||
manager = qiniu.BucketManager(auth)
|
||
marker: str | None = None
|
||
deleted = 0
|
||
while True:
|
||
result, eof, info = manager.list(settings.qiniu_bucket, prefix=prefix, marker=marker, limit=1000)
|
||
if info.status_code != 200:
|
||
raise RuntimeError(f"七牛列举失败：{info.error or info.text_body or info.status_code}")
|
||
items = (result or {}).get("items") or []
|
||
for item in items:
|
||
await self.delete(item["key"])
|
||
deleted += 1
|
||
marker = (result or {}).get("marker")
|
||
if eof or not marker:
|
||
return deleted
|
||
|
||
|
||
def get_storage():
|
||
settings = get_settings()
|
||
if settings.use_qiniu:
|
||
return QiniuStorage()
|
||
return LocalStorage()
|
||
|
||
|
||
def storage_key_from_url(url: str) -> str | None:
|
||
"""将本项目生成的公开 URL 还原为对象 key；第三方源站 URL 返回 None。"""
|
||
if not url:
|
||
return None
|
||
settings = get_settings()
|
||
if "/media/" in url:
|
||
return url.split("/media/", 1)[1].split("?", 1)[0]
|
||
domain = settings.qiniu_domain.rstrip("/")
|
||
if domain and url.startswith(f"{domain}/"):
|
||
return url[len(domain) + 1:].split("?", 1)[0]
|
||
return None
|
||
|
||
|
||
def local_path(stored_url_or_key: str) -> Path | None:
|
||
"""stored_url（http…/media/xxx）或 media key → 本地文件路径。
|
||
|
||
仅本地存储模式下可命中；七牛 URL 不含 /media/，返回 None 由调用方回退远程下载。
|
||
"""
|
||
s = stored_url_or_key or ""
|
||
if "/media/" in s:
|
||
s = s.split("/media/", 1)[1]
|
||
root = _LOCAL_ROOT.resolve()
|
||
p = (root / s.split("?", 1)[0]).resolve()
|
||
# URL / key 可能来自历史数据，不能让 ../../ 逃出本地媒体目录。
|
||
if root not in p.parents:
|
||
return None
|
||
return p if p.is_file() else None
|