SFM/app/pipeline/attachments.py

"""Сбор XML-файлов из вложений интерфейса.

Пользователь прикрепляет через file_uploader сразу несколько объектов:
отдельные .xml, архивы с выгрузками, или и то и другое вперемешку. Этот
модуль приводит их к плоскому списку XML-файлов на диске и отдельно
возвращает список вложений, которые обработать не удалось, — чтобы одно
битое вложение не отменяло весь запуск.

Формат архива определяется по содержимому (zipfile/tarfile), а не по
расширению: выгрузки нередко приезжают с именем вида "Файлы" без суффикса
или с суффиксом, не совпадающим с фактическим форматом.
"""

from __future__ import annotations

import tarfile
import zipfile
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path

XML_SUFFIX = ".xml"
# Форматы, для которых нужен внешний бинарник или сторонняя зависимость,
# поэтому распознаём их только чтобы выдать понятную причину отказа.
UNSUPPORTED_ARCHIVE_SUFFIXES = frozenset({".rar", ".7z"})


class AttachmentError(Exception):
    """Вложение не удалось разобрать: не архив, битый архив, опасные пути."""


@dataclass(frozen=True)
class CollectedXml:
    """XML-файл, готовый к разбору, и вложение, из которого он получен."""

    file_name: str
    path: Path
    source: str


@dataclass(frozen=True)
class AttachmentProblem:
    source: str
    reason: str
    # Для дубликатов — имя самого файла; для нечитаемых вложений пусто,
    # потому что до отдельных файлов дело не доходит.
    file_name: str = ""


@dataclass
class CollectResult:
    files: list[CollectedXml] = field(default_factory=list)
    problems: list[AttachmentProblem] = field(default_factory=list)
    duplicates: list[AttachmentProblem] = field(default_factory=list)
    archives: int = 0
    loose_xml_files: int = 0


def collect_xml_attachments(
    attachments: Sequence[Path],
    extract_root: Path,
) -> CollectResult:
    """Разворачивает вложения в плоский список XML-файлов.

    Дубликаты по имени файла отбрасываются: имя XML-выгрузки уникально
    (в нём есть дата и порядковый номер), поэтому повторное имя означает,
    что одна и та же выгрузка прикреплена дважды, а её повторный разбор
    удвоил бы строки отчёта.
    """
    result = CollectResult()
    seen: dict[str, str] = {}

    for index, attachment in enumerate(attachments):
        source = attachment.name
        if attachment.suffix.lower() == XML_SUFFIX:
            result.loose_xml_files += 1
            _add_file(result, seen, file_name=source, path=attachment, source=source)
            continue

        # Каждое вложение распаковываем в свой каталог: имена файлов внутри
        # разных архивов могут совпадать.
        target_dir = extract_root / str(index)
        try:
            target_dir.mkdir(parents=True, exist_ok=True)
            extract_archive(attachment, target_dir)
        except AttachmentError as exc:
            result.problems.append(AttachmentProblem(source=source, reason=str(exc)))
            continue

        result.archives += 1
        for xml_path in sorted(target_dir.rglob("*")):
            if xml_path.is_file() and xml_path.suffix.lower() == XML_SUFFIX:
                _add_file(
                    result,
                    seen,
                    file_name=xml_path.name,
                    path=xml_path,
                    source=source,
                )

    return result


def extract_archive(archive_path: Path, target_dir: Path) -> None:
    """Распаковывает zip или tar в target_dir, не выходя за его пределы."""
    if zipfile.is_zipfile(archive_path):
        _extract_zip(archive_path, target_dir)
        return
    if tarfile.is_tarfile(archive_path):
        _extract_tar(archive_path, target_dir)
        return

    suffix = archive_path.suffix.lower()
    if suffix in UNSUPPORTED_ARCHIVE_SUFFIXES:
        raise AttachmentError(
            f"формат {suffix} не поддерживается — прикрепите .zip или .tar"
        )
    raise AttachmentError("вложение не является ни XML-файлом, ни архивом .zip/.tar")


def _extract_zip(archive_path: Path, target_dir: Path) -> None:
    try:
        with zipfile.ZipFile(archive_path) as archive:
            for member in archive.infolist():
                if member.is_dir():
                    continue
                destination = _safe_destination(target_dir, member.filename)
                destination.parent.mkdir(parents=True, exist_ok=True)
                with archive.open(member) as source, destination.open("wb") as target:
                    target.write(source.read())
    except (zipfile.BadZipFile, OSError) as exc:
        raise AttachmentError(f"не удалось распаковать zip-архив: {exc}") from exc


def _extract_tar(archive_path: Path, target_dir: Path) -> None:
    # Своя распаковка вместо extractall: параметр filter="data", который
    # отсекает опасных членов архива, появился только в Python 3.11.4, а
    # окружение бывает и старше.
    try:
        with tarfile.open(archive_path) as archive:
            for member in archive.getmembers():
                if member.isdir():
                    continue
                if not member.isfile():
                    # Симлинки и устройства из архива нам не нужны и опасны:
                    # ссылка может указать за пределы каталога распаковки.
                    continue
                destination = _safe_destination(target_dir, member.name)
                destination.parent.mkdir(parents=True, exist_ok=True)
                extracted = archive.extractfile(member)
                if extracted is None:
                    continue
                with extracted, destination.open("wb") as target:
                    target.write(extracted.read())
    except (tarfile.TarError, OSError) as exc:
        raise AttachmentError(f"не удалось распаковать tar-архив: {exc}") from exc


def _safe_destination(target_dir: Path, member_name: str) -> Path:
    """Защита от путей вида "../..": член архива обязан лежать внутри каталога."""
    base = target_dir.resolve()
    # PurePath не трогает разделители Windows, а в архивах они бывают оба.
    relative_parts = [
        part
        for part in member_name.replace("\\", "/").split("/")
        if part and part not in {".", ".."}
    ]
    if not relative_parts:
        raise AttachmentError(f"архив содержит недопустимый путь: {member_name!r}")
    destination = base.joinpath(*relative_parts)
    if not destination.is_relative_to(base):
        raise AttachmentError(f"архив содержит недопустимый путь: {member_name!r}")
    return destination


def _add_file(
    result: CollectResult,
    seen: dict[str, str],
    *,
    file_name: str,
    path: Path,
    source: str,
) -> None:
    previous_source = seen.get(file_name)
    if previous_source is not None:
        result.duplicates.append(
            AttachmentProblem(
                source=source,
                file_name=file_name,
                reason=(
                    f"Повторно прикреплённый файл — уже получен из вложения "
                    f"«{previous_source}», пропущен."
                ),
            )
        )
        return
    seen[file_name] = source
    result.files.append(CollectedXml(file_name=file_name, path=path, source=source))