100 lines
4.8 KiB
Python
100 lines
4.8 KiB
Python
"""Bounded ZIP extraction for packages uploaded to the AI Core host."""
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import re
|
|
import shutil
|
|
import stat
|
|
import tempfile
|
|
import zipfile
|
|
import zlib
|
|
from pathlib import Path
|
|
from collections.abc import Callable
|
|
from typing import TypeVar
|
|
|
|
from app.errors import ApiError
|
|
from app.extensions.errors import ExtensionError
|
|
|
|
MAX_ZIP_BYTES = 10 * 1024 * 1024
|
|
MAX_EXPANDED_BYTES = 50 * 1024 * 1024
|
|
MAX_ENTRIES = 2048
|
|
T = TypeVar('T')
|
|
|
|
|
|
def invalid(message: str) -> ApiError:
|
|
return ApiError(422, 'EXTENSION_ZIP_INVALID', message)
|
|
|
|
|
|
def install_zip(data: bytes, kind: str, storage: Path, install: Callable[[Path], T], *, managed_install: Callable[[Path, Path], T] | None = None) -> T:
|
|
if len(data) > MAX_ZIP_BYTES:
|
|
raise ApiError(413, 'EXTENSION_ZIP_TOO_LARGE', 'ZIP 文件不能超过 10 MiB。')
|
|
if kind not in ('skill', 'plugin'):
|
|
raise ValueError('Unknown extension kind')
|
|
storage.mkdir(parents=True, exist_ok=True)
|
|
# Retain successful extraction: Plugin commands and resources use this directory.
|
|
destination = Path(tempfile.mkdtemp(prefix=f'{kind}-', dir=storage))
|
|
try:
|
|
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
|
entries = archive.infolist()
|
|
if not entries or len(entries) > MAX_ENTRIES:
|
|
raise invalid('ZIP 为空或文件条目超过 2048 个。')
|
|
seen: set[str] = set()
|
|
spellings: dict[str, str] = {}
|
|
total = 0
|
|
for entry in entries:
|
|
name = entry.filename.rstrip('/')
|
|
parts = name.split('/')
|
|
if (entry.orig_filename != entry.filename or '\\' in name
|
|
or any(not p or p in ('.', '..') or any(c in p for c in ':*?<>|"') or p.endswith((' ', '.'))
|
|
or any(ord(c) < 32 for c in p)
|
|
or re.match(r'^(CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\.|$)', p, re.I)
|
|
for p in parts)):
|
|
raise invalid('ZIP 包含不安全的文件路径。')
|
|
mode = stat.S_IFMT(entry.external_attr >> 16)
|
|
if mode not in (0, stat.S_IFREG, stat.S_IFDIR) or entry.flag_bits & 1:
|
|
raise invalid('ZIP 不支持链接、特殊文件或加密条目。')
|
|
if entry.compress_type not in (zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED):
|
|
raise invalid('ZIP 仅支持 stored/deflate 压缩。')
|
|
key = name.casefold()
|
|
if key in seen:
|
|
raise invalid('ZIP 包含重复或大小写冲突的路径。')
|
|
seen.add(key)
|
|
for index in range(1, len(parts) + 1):
|
|
prefix = '/'.join(parts[:index])
|
|
if spellings.setdefault(prefix.casefold(), prefix) != prefix:
|
|
raise invalid('ZIP 包含大小写冲突的目录。')
|
|
total += entry.file_size
|
|
if total > MAX_EXPANDED_BYTES:
|
|
raise ApiError(413, 'EXTENSION_ZIP_TOO_LARGE', 'ZIP 解压后不能超过 50 MiB。')
|
|
target = destination.joinpath(*parts)
|
|
if not target.resolve().is_relative_to(destination.resolve()):
|
|
raise invalid('ZIP 路径超出包目录。')
|
|
written = 0
|
|
for entry in entries:
|
|
target = destination.joinpath(*entry.filename.rstrip('/').split('/'))
|
|
if entry.is_dir():
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
continue
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
with archive.open(entry) as source, target.open('xb') as output:
|
|
while chunk := source.read(64 * 1024):
|
|
written += len(chunk)
|
|
if written > MAX_EXPANDED_BYTES:
|
|
raise ApiError(413, 'EXTENSION_ZIP_TOO_LARGE', 'ZIP 解压后不能超过 50 MiB。')
|
|
output.write(chunk)
|
|
manifest = f'{kind}.yaml'
|
|
root = destination
|
|
if not (root / manifest).is_file():
|
|
children = list(root.iterdir())
|
|
if len(children) != 1 or not children[0].is_dir() or not (children[0] / manifest).is_file():
|
|
raise invalid(f'ZIP 根目录或唯一顶层文件夹中须包含 {manifest}。')
|
|
root = children[0]
|
|
return managed_install(root, destination) if managed_install else install(root)
|
|
except BaseException as error:
|
|
shutil.rmtree(destination)
|
|
if isinstance(error, ExtensionError):
|
|
raise
|
|
if isinstance(error, (zipfile.BadZipFile, OSError, RuntimeError, NotImplementedError, zlib.error, EOFError, UnicodeError)):
|
|
raise invalid('ZIP 损坏、路径冲突或无法解压。') from error
|
|
raise
|