216 lines
8.1 KiB
Python
216 lines
8.1 KiB
Python
"""对象存储上传工具（MinIO，S3 协议兼容）
|
||
|
||
原实现基于阿里云 OSS SDK（alibabacloud_oss_v2），现改为 boto3 对接 MinIO。
|
||
对外函数名与返回值保持不变，业务代码无需修改。
|
||
"""
|
||
|
||
import base64
|
||
import io
|
||
import mimetypes
|
||
import os
|
||
import re
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
|
||
import boto3
|
||
import requests
|
||
from botocore.config import Config
|
||
|
||
from config import (
|
||
region,
|
||
endpoint,
|
||
bucket,
|
||
file_url_pre,
|
||
bucket_path,
|
||
accessKeyId,
|
||
accessKeySecret,
|
||
)
|
||
from utils.ssrf import is_internal_url
|
||
|
||
_client = None
|
||
_client_lock = threading.Lock()
|
||
_REMOTE_IMAGE_MAX_BYTES = int(os.getenv("OSS_UPLOAD_IMAGE_MAX_BYTES", str(10 * 1024 * 1024)))
|
||
_REMOTE_IMAGE_TIMEOUT = (5, 30)
|
||
|
||
|
||
def get_client():
|
||
"""获取（懒加载）S3 客户端，MinIO 必须使用 path-style 寻址"""
|
||
global _client
|
||
if _client is None:
|
||
with _client_lock:
|
||
if _client is None:
|
||
_client = boto3.client(
|
||
"s3",
|
||
endpoint_url=endpoint,
|
||
aws_access_key_id=accessKeyId,
|
||
aws_secret_access_key=accessKeySecret,
|
||
region_name=region,
|
||
config=Config(
|
||
signature_version="s3v4",
|
||
s3={"addressing_style": "path"}, # MinIO 必须
|
||
retries={"max_attempts": 3, "mode": "standard"},
|
||
),
|
||
)
|
||
return _client
|
||
|
||
|
||
def _guess_content_type(key: str) -> str:
|
||
content_type, _ = mimetypes.guess_type(key)
|
||
return content_type or "application/octet-stream"
|
||
|
||
|
||
def get_presigned_url(key: str, expires: int = 7 * 24 * 3600) -> str:
|
||
"""生成临时访问链接（存储桶未开放匿名读时使用），默认有效 7 天"""
|
||
return get_client().generate_presigned_url(
|
||
"get_object",
|
||
Params={"Bucket": bucket, "Key": key.lstrip("/")},
|
||
ExpiresIn=expires,
|
||
)
|
||
|
||
|
||
|
||
class _LimitedReader:
|
||
"""Keep a streamed upload from accepting an unexpectedly huge object."""
|
||
|
||
def __init__(self, source, max_bytes: int):
|
||
self._source = source
|
||
self._max_bytes = max_bytes
|
||
self._read_bytes = 0
|
||
|
||
def read(self, size=-1):
|
||
remaining = self._max_bytes - self._read_bytes
|
||
if remaining < 0:
|
||
raise ValueError("upload exceeds configured size limit")
|
||
read_size = remaining + 1 if size is None or size < 0 else min(size, remaining + 1)
|
||
data = self._source.read(read_size)
|
||
if not data:
|
||
return data
|
||
self._read_bytes += len(data)
|
||
if self._read_bytes > self._max_bytes:
|
||
raise ValueError("upload exceeds configured size limit")
|
||
return data
|
||
|
||
def seek(self, offset, whence=0):
|
||
position = self._source.seek(offset, whence)
|
||
self._read_bytes = max(0, position)
|
||
return position
|
||
|
||
def tell(self):
|
||
return self._source.tell()
|
||
|
||
|
||
def upload_fileobj(file_obj, key: str, max_bytes: int = 0, content_type: str = ""):
|
||
"""Stream a file-like object to S3/MinIO without materializing it."""
|
||
key = key.lstrip("/")
|
||
if isinstance(file_obj, (bytes, bytearray)):
|
||
file_obj = io.BytesIO(file_obj)
|
||
if not hasattr(file_obj, "read"):
|
||
raise TypeError("file_obj must be bytes or a readable file-like object")
|
||
stream_size = None
|
||
try:
|
||
file_obj.seek(0)
|
||
file_obj.seek(0, os.SEEK_END)
|
||
stream_size = file_obj.tell()
|
||
file_obj.seek(0)
|
||
except (AttributeError, OSError, TypeError, ValueError):
|
||
try:
|
||
file_obj.seek(0)
|
||
except (AttributeError, OSError, TypeError, ValueError):
|
||
pass
|
||
if max_bytes and max_bytes > 0 and stream_size is not None and stream_size > max_bytes:
|
||
raise ValueError("upload exceeds configured size limit")
|
||
client = get_client()
|
||
body = _LimitedReader(file_obj, max_bytes) if max_bytes and max_bytes > 0 else file_obj
|
||
client.upload_fileobj(
|
||
body,
|
||
bucket,
|
||
key,
|
||
ExtraArgs={"ContentType": content_type or _guess_content_type(key)},
|
||
)
|
||
return file_url_pre + key
|
||
|
||
|
||
def upload_file(file_content: bytes, key: str):
|
||
"""上传字节内容到 MinIO，返回可访问链接"""
|
||
return upload_fileobj(file_content, key)
|
||
|
||
|
||
def upload_data_url(data_url: str, prefix: str = "history", key_hint: str = "") -> str:
|
||
"""
|
||
将 base64 data URL 上传到对象存储，返回图片链接
|
||
data_url: data:image/png;base64,xxxx 或 data:image/jpeg;base64,xxxx
|
||
prefix: 对象 key 前缀
|
||
key_hint: 可选后缀避免重名，如 "_0", "_1"
|
||
"""
|
||
match = re.match(r'data:image/(\w+);base64,(.+)', data_url)
|
||
if not match:
|
||
raise ValueError('无效的 data URL 格式')
|
||
ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
|
||
encoded_payload = match.group(2)
|
||
if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
file_content = base64.b64decode(encoded_payload, validate=True)
|
||
if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
ts = int(time.time() * 1000)
|
||
key = f"{bucket_path}{prefix}/{ts}{key_hint}.{ext}"
|
||
return upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES)
|
||
|
||
|
||
def upload_data_urls(data_urls: list, prefix: str = "history") -> list:
|
||
"""批量上传图片；每张图���完成上传后立即释放其缓冲区。"""
|
||
urls = []
|
||
ts = int(time.time() * 1000)
|
||
for i, data_url in enumerate(data_urls or []):
|
||
if not data_url or not isinstance(data_url, str):
|
||
continue
|
||
if not data_url.startswith("http"):
|
||
match = re.match(r'data:image/(\w+);base64,(.+)', data_url)
|
||
if not match:
|
||
continue
|
||
ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
|
||
encoded_payload = match.group(2)
|
||
if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
file_content = base64.b64decode(encoded_payload, validate=True)
|
||
if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
|
||
urls.append(upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES))
|
||
del file_content
|
||
continue
|
||
|
||
if is_internal_url(data_url):
|
||
raise ValueError("拒绝下载内网/本机地址的图片")
|
||
|
||
with requests.get(data_url, stream=True, timeout=_REMOTE_IMAGE_TIMEOUT) as response:
|
||
response.raise_for_status()
|
||
content_length = response.headers.get("Content-Length")
|
||
if content_length and int(content_length) > _REMOTE_IMAGE_MAX_BYTES:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
|
||
extension = mimetypes.guess_extension(content_type) or ".png"
|
||
ext = extension.lstrip(".") or "png"
|
||
key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
|
||
with tempfile.SpooledTemporaryFile(max_size=2 * 1024 * 1024, mode="w+b") as image_file:
|
||
total = 0
|
||
for chunk in response.iter_content(chunk_size=64 * 1024):
|
||
if not chunk:
|
||
continue
|
||
total += len(chunk)
|
||
if total > _REMOTE_IMAGE_MAX_BYTES:
|
||
raise ValueError("image exceeds configured upload size limit")
|
||
image_file.write(chunk)
|
||
image_file.seek(0)
|
||
urls.append(upload_fileobj(image_file, key, _REMOTE_IMAGE_MAX_BYTES, content_type))
|
||
return urls
|
||
|
||
|
||
# 脚本入口，当文件被直接运行时调用main函数
|
||
if __name__ == "__main__":
|
||
with open("D:\\pack\\nanri\\main.dist\\SHUFU.zip", "rb") as f:
|
||
file_content = f.read()
|
||
res = upload_file(file_content, key=bucket_path + "versions/1.0.46.zip")
|
||
print(res)
|