100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
import hashlib
|
|
import logging
|
|
import mimetypes
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import HTTPException, Response
|
|
from .config import settings
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _safe_extension(filename: str, content_type: str) -> str:
|
|
extension = Path(filename or "").suffix.lower()
|
|
if extension in {".jpg", ".jpeg", ".png", ".webp"}:
|
|
return extension
|
|
guessed = mimetypes.guess_extension(content_type or "") or ".bin"
|
|
return ".jpg" if guessed == ".jpe" else guessed
|
|
|
|
|
|
def _absolute_path(relative_path: str) -> Path:
|
|
root = settings.screenshot_storage_dir.resolve()
|
|
candidate = (root / relative_path).resolve()
|
|
if root != candidate and root not in candidate.parents:
|
|
raise ValueError("非法截图文件路径")
|
|
return candidate
|
|
|
|
|
|
def store_screenshot(data: bytes, filename: str, content_type: str, category: str) -> str:
|
|
if not data:
|
|
raise ValueError("截图内容不能为空")
|
|
now_path = Path(category) / uuid.uuid4().hex[:2]
|
|
relative_path = now_path / f"{uuid.uuid4().hex}{_safe_extension(filename, content_type)}"
|
|
target = _absolute_path(relative_path.as_posix())
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = target.with_suffix(target.suffix + ".tmp")
|
|
temporary.write_bytes(data)
|
|
if temporary.stat().st_size != len(data):
|
|
temporary.unlink(missing_ok=True)
|
|
raise OSError("截图写入不完整")
|
|
temporary.replace(target)
|
|
return relative_path.as_posix()
|
|
|
|
|
|
def screenshot_fields(data: bytes, filename: str, content_type: str, category: str) -> dict:
|
|
try:
|
|
file_path = store_screenshot(data, filename, content_type, category)
|
|
except (OSError, ValueError) as exc:
|
|
logger.exception(
|
|
"Screenshot file write failed: category=%s filename=%s storage_dir=%s",
|
|
category,
|
|
filename,
|
|
settings.screenshot_storage_dir,
|
|
)
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="图片存储目录不可用，请联系管理员检查目录配置和写入权限",
|
|
) from exc
|
|
return {
|
|
"filename": filename,
|
|
"content_type": content_type,
|
|
"file_path": file_path,
|
|
"image_data": None,
|
|
}
|
|
|
|
|
|
def screenshot_bytes(image) -> bytes:
|
|
if image.file_path:
|
|
try:
|
|
path = _absolute_path(image.file_path)
|
|
if path.is_file():
|
|
return path.read_bytes()
|
|
except (OSError, ValueError):
|
|
pass
|
|
if image.image_data:
|
|
return bytes(image.image_data)
|
|
raise HTTPException(status_code=404, detail="截图文件不存在")
|
|
|
|
|
|
def screenshot_response(image) -> Response:
|
|
return Response(
|
|
content=screenshot_bytes(image),
|
|
media_type=image.content_type or "application/octet-stream",
|
|
headers={"Cache-Control": "private, max-age=600"},
|
|
)
|
|
|
|
|
|
def screenshot_data_url(image) -> str:
|
|
import base64
|
|
|
|
encoded = base64.b64encode(screenshot_bytes(image)).decode()
|
|
return f"data:{image.content_type};base64,{encoded}"
|
|
|
|
|
|
def file_digest(relative_path: str) -> tuple[int, str]:
|
|
data = _absolute_path(relative_path).read_bytes()
|
|
return len(data), hashlib.sha256(data).hexdigest()
|
|
|