"""文件键由服务端生成,扩展名不能绕过类型验证;CPU 转图在后台线程运行。""" import base64 from io import BytesIO from pathlib import Path import uuid import pymupdf as fitz from PIL import Image from app.shared.common import BusinessError from app.shared.config import settings Image.MAX_IMAGE_PIXELS = 20_000_000 def validate_file(data: bytes) -> tuple[str, str]: """上传即验证页数和图像解码,不能把任意二进制默认当 JPEG。""" if not data or len(data) > settings.max_upload_mb * 1024 * 1024: raise BusinessError("文件为空或超过上传大小上限", code="UPLOAD_SIZE") try: if data.startswith(b"%PDF"): with fitz.open(stream=data, filetype="pdf") as doc: if doc.is_encrypted or not 1 <= len(doc) <= settings.max_pdf_pages: raise BusinessError("PDF 加密或页数超过限制", code="PDF_LIMIT") return "application/pdf", ".pdf" if data.startswith(b"\x89PNG\r\n\x1a\n") or data.startswith(b"\xff\xd8\xff"): with Image.open(BytesIO(data)) as image: image.verify() if image.width * image.height > Image.MAX_IMAGE_PIXELS: raise BusinessError("图像像素超过上限", code="IMAGE_LIMIT") return ("image/png", ".png") if data.startswith(b"\x89PNG") else ("image/jpeg", ".jpg") except BusinessError: raise except Exception: raise BusinessError("文件损坏或无法解析", code="FILE_INVALID") raise BusinessError("只接受真实 PNG、JPEG 和 PDF 文件", code="FILE_TYPE") def save_file(data: bytes, extension: str) -> str: key = uuid.uuid4().hex + extension folder = settings.data_path / "uploads" folder.mkdir(parents=True, exist_ok=True) (folder / key).write_bytes(data) return key def file_path(key: str) -> Path: """从数据库取受控键,额外边界检查防止损坏数据库造成路径穿越。""" folder = (settings.data_path / "uploads").resolve() target = (folder / key).resolve() if target.parent != folder or not target.is_file(): raise BusinessError("发票文件不存在", 404, "FILE_NOT_FOUND") return target def vision_images(key: str, content_type: str) -> list[str]: """PDF 逐页转图而非把 PDF 当 image_url;限制页数保证不会静默漏页。""" path = file_path(key) images = [] if content_type == "application/pdf": with fitz.open(path) as doc: if len(doc) > settings.max_pdf_pages: raise BusinessError("PDF 超过允许页数", code="PDF_LIMIT") for page in doc: scale = min(2.0, 1536 / max(page.rect.width, page.rect.height)) pixmap = page.get_pixmap(matrix=fitz.Matrix(scale, scale), alpha=False) images.append(Image.open(BytesIO(pixmap.tobytes("png"))).copy()) else: with Image.open(path) as image: images.append(image.convert("RGB")) result = [] for image in images: image.thumbnail((1536, 1536)) output = BytesIO() image.convert("RGB").save(output, format="JPEG", quality=90) result.append("data:image/jpeg;base64," + base64.b64encode(output.getvalue()).decode()) return result