R524809
433 lines
17 KiB
Python
433 lines
17 KiB
Python
"""图片/文件存储抽象：本地文件系统（开发兜底）+ 七牛（生产）。"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import io
|
||
import ipaddress
|
||
import mimetypes
|
||
import socket
|
||
import shutil
|
||
import uuid
|
||
import warnings
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from urllib.parse import urljoin, urlparse
|
||
|
||
import httpx
|
||
from PIL import Image, ImageOps, UnidentifiedImageError
|
||
|
||
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)
|
||
|
||
|
||
class UnsafeRemoteUrlError(ValueError):
|
||
"""远程资源 URL 指向回环、内网或其他保留网络，不能由服务端代请求。"""
|
||
|
||
|
||
class ResourceDownloadTimeoutError(TimeoutError):
|
||
"""远程资源下载超过连接、读取或总耗时限制。"""
|
||
|
||
|
||
class ImageValidationError(ValueError):
|
||
"""图片字节无法解码，或不符合格式与像素限制。"""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class NormalizedImage:
|
||
"""经真实格式校验后的图片字节及可落库元数据。"""
|
||
|
||
data: bytes
|
||
content_type: str
|
||
width: int
|
||
height: int
|
||
image_format: str
|
||
orientation_normalized: bool
|
||
|
||
|
||
_IMAGE_CONTENT_TYPES = {
|
||
"JPEG": "image/jpeg",
|
||
"PNG": "image/png",
|
||
"WEBP": "image/webp",
|
||
"GIF": "image/gif",
|
||
"BMP": "image/bmp",
|
||
}
|
||
|
||
|
||
def normalize_image_bytes(
|
||
data: bytes,
|
||
max_bytes: int | None = None,
|
||
max_pixels: int | None = None,
|
||
) -> NormalizedImage:
|
||
"""按真实图片内容校验、纠正 EXIF 方向，并返回可信 MIME 与宽高。
|
||
|
||
不盲从 URL 后缀或客户端 Content-Type。为避免无意义的质量损耗，仅当 EXIF
|
||
方向不是默认值时才重编码；动画图不重编码，以免丢失帧。
|
||
"""
|
||
settings = get_settings()
|
||
byte_limit = max_bytes if max_bytes is not None else settings.image_upload_max_bytes
|
||
pixel_limit = max_pixels if max_pixels is not None else settings.image_max_pixels
|
||
if byte_limit <= 0 or pixel_limit <= 0:
|
||
raise ValueError("图片大小和像素上限必须大于 0")
|
||
if not data:
|
||
raise ImageValidationError("图片内容为空")
|
||
if len(data) > byte_limit:
|
||
raise ResourceSizeLimitError(byte_limit, len(data))
|
||
|
||
try:
|
||
# Pillow 的默认 bomb 保护不能替代项目自定义像素上限；把 warning 升级为错误。
|
||
with warnings.catch_warnings():
|
||
warnings.simplefilter("error", Image.DecompressionBombWarning)
|
||
with Image.open(io.BytesIO(data)) as image:
|
||
image_format = (image.format or "").upper()
|
||
if image_format not in _IMAGE_CONTENT_TYPES:
|
||
raise ImageValidationError(f"不支持的图片格式：{image_format or '未知'}")
|
||
|
||
width, height = image.size
|
||
if width <= 0 or height <= 0 or width * height > pixel_limit:
|
||
raise ImageValidationError(
|
||
f"图片像素超过上限（最多 {pixel_limit:,}，当前 {width:,} × {height:,}）"
|
||
)
|
||
|
||
# 强制实际解码，拒绝截断/伪造图片而不是只相信文件头。
|
||
image.load()
|
||
frame_count = getattr(image, "n_frames", 1)
|
||
orientation = image.getexif().get(274, 1) if frame_count == 1 else 1
|
||
normalized = orientation not in (None, 1)
|
||
output = data
|
||
if normalized:
|
||
corrected = ImageOps.exif_transpose(image)
|
||
width, height = corrected.size
|
||
if image_format == "JPEG" and corrected.mode not in {"RGB", "L"}:
|
||
corrected = corrected.convert("RGB")
|
||
buffer = io.BytesIO()
|
||
corrected.save(buffer, format=image_format)
|
||
output = buffer.getvalue()
|
||
except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc:
|
||
raise ImageValidationError("文件不是可用的图片，或图片数据已损坏") from exc
|
||
except Image.DecompressionBombWarning as exc:
|
||
raise ImageValidationError("图片像素过大，已拒绝解码") from exc
|
||
|
||
if len(output) > byte_limit:
|
||
raise ResourceSizeLimitError(byte_limit, len(output))
|
||
return NormalizedImage(
|
||
data=output,
|
||
content_type=_IMAGE_CONTENT_TYPES[image_format],
|
||
width=width,
|
||
height=height,
|
||
image_format=image_format,
|
||
orientation_normalized=normalized,
|
||
)
|
||
|
||
|
||
def _is_public_ip(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||
"""仅允许可公开路由的地址；IPv4-mapped IPv6 需按其 IPv4 实际地址判断。"""
|
||
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped:
|
||
return _is_public_ip(address.ipv4_mapped)
|
||
return address.is_global
|
||
|
||
|
||
async def validate_remote_url(url: str) -> None:
|
||
"""校验服务端允许代请求的公网 HTTP(S) URL。
|
||
|
||
该校验用于代理、采集转存、AI 参考图和导出。项目自己的 `/media/` 文件
|
||
在调用方会先直读磁盘；其他回环或私网 URL 一律不能通过网络访问。
|
||
"""
|
||
try:
|
||
parsed = urlparse(url)
|
||
port = parsed.port
|
||
except ValueError as exc:
|
||
raise UnsafeRemoteUrlError("图片地址端口无效") from exc
|
||
|
||
if parsed.scheme.lower() not in {"http", "https"}:
|
||
raise UnsafeRemoteUrlError("仅允许 http/https 图片地址")
|
||
if not parsed.hostname:
|
||
raise UnsafeRemoteUrlError("图片地址缺少主机名")
|
||
if parsed.username is not None or parsed.password is not None:
|
||
raise UnsafeRemoteUrlError("图片地址不能包含用户名或密码")
|
||
|
||
host = parsed.hostname.rstrip(".")
|
||
if host.lower() in {"localhost", "localhost.localdomain"}:
|
||
raise UnsafeRemoteUrlError("不允许请求 localhost 或本机网络地址")
|
||
|
||
try:
|
||
addresses = {ipaddress.ip_address(host)}
|
||
except ValueError:
|
||
try:
|
||
infos = await asyncio.get_running_loop().getaddrinfo(
|
||
host,
|
||
port or (443 if parsed.scheme.lower() == "https" else 80),
|
||
type=socket.SOCK_STREAM,
|
||
)
|
||
except socket.gaierror as exc:
|
||
raise UnsafeRemoteUrlError("图片地址域名无法解析") from exc
|
||
addresses = {ipaddress.ip_address(info[4][0]) for info in infos}
|
||
|
||
if not addresses:
|
||
raise UnsafeRemoteUrlError("图片地址域名未解析到有效地址")
|
||
private_addresses = [str(address) for address in addresses if not _is_public_ip(address)]
|
||
if private_addresses:
|
||
raise UnsafeRemoteUrlError(
|
||
f"不允许请求非公网地址（{', '.join(sorted(private_addresses))}）"
|
||
)
|
||
|
||
|
||
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"
|
||
if data[:2] == b"\xff\xd8":
|
||
return ".jpg"
|
||
if data[:2] == b"BM":
|
||
return ".bmp"
|
||
return _ext_from_content_type(content_type)
|
||
|
||
|
||
async def download_bytes(
|
||
url: str,
|
||
referer: str | None = None,
|
||
timeout: float | None = None,
|
||
max_bytes: int | None = None,
|
||
) -> tuple[bytes, str]:
|
||
"""安全地流式下载公网资源，并限制地址、重定向、体积和耗时。"""
|
||
settings = get_settings()
|
||
limit = max_bytes or settings.image_download_max_bytes
|
||
if limit <= 0:
|
||
raise ValueError("图片下载大小上限必须大于 0")
|
||
total_timeout = timeout if timeout is not None else settings.image_download_total_timeout
|
||
if total_timeout <= 0:
|
||
raise ValueError("图片下载总超时必须大于 0")
|
||
if settings.image_download_connect_timeout <= 0 or settings.image_download_read_timeout <= 0:
|
||
raise ValueError("图片下载连接和读取超时必须大于 0")
|
||
if settings.image_download_max_redirects < 0:
|
||
raise ValueError("图片下载重定向上限不能小于 0")
|
||
|
||
headers = {"Referer": referer} if referer else {}
|
||
http_timeout = httpx.Timeout(
|
||
connect=settings.image_download_connect_timeout,
|
||
read=settings.image_download_read_timeout,
|
||
write=settings.image_download_read_timeout,
|
||
pool=settings.image_download_connect_timeout,
|
||
)
|
||
current_url = url
|
||
try:
|
||
async with asyncio.timeout(total_timeout):
|
||
async with httpx.AsyncClient(timeout=http_timeout, follow_redirects=False) as client:
|
||
for redirect_count in range(settings.image_download_max_redirects + 1):
|
||
await validate_remote_url(current_url)
|
||
async with client.stream("GET", current_url, headers=headers) as resp:
|
||
if resp.is_redirect:
|
||
location = resp.headers.get("location")
|
||
if not location:
|
||
raise httpx.RemoteProtocolError("重定向响应缺少 Location")
|
||
if redirect_count >= settings.image_download_max_redirects:
|
||
raise httpx.TooManyRedirects("图片重定向次数超过上限")
|
||
current_url = urljoin(str(resp.url), location)
|
||
continue
|
||
|
||
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
|
||
except (TimeoutError, httpx.TimeoutException) as exc:
|
||
raise ResourceDownloadTimeoutError(f"图片下载超时（最多 {total_timeout:g} 秒）") from exc
|
||
|
||
raise RuntimeError("图片下载未返回有效响应")
|
||
|
||
|
||
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
|