code/ai_finance_agent/backend/app/infrastructure/files.py

"""文件键由服务端生成，扩展名不能绕过类型验证；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