wyndham-ARR/arr_web/downloads.py

"""Private artifact download ports; object paths never enter public JSON."""

from __future__ import annotations

import hashlib
import os
import stat
import tempfile
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Optional, Protocol

from arr_ingestion.contracts import IngestionError
from arr_storage.store import ManagedObjectStore
from arr_web.contracts import PortalError


MAX_DOWNLOAD_BYTES = 100 * 1024 * 1024
ALLOWED_DOWNLOAD_KINDS = frozenset(
    {"daily_xlsx", "monthly_xlsx", "company_ten_day_xlsx"}
)


@dataclass(frozen=True)
class ArtifactDescriptor:
    file_kind: str
    original_filename: str
    storage_key: str
    sha256: str
    byte_size: int
    mime_type: str

    def validate(self) -> None:
        path = PurePosixPath(self.storage_key)
        if (
            self.file_kind not in ALLOWED_DOWNLOAD_KINDS
            or not self.original_filename
            or PurePosixPath(self.original_filename).name != self.original_filename
            or "/" in self.original_filename
            or "\\" in self.original_filename
            or "\x00" in self.original_filename
            or not self.storage_key
            or path.is_absolute()
            or "." in path.parts
            or ".." in path.parts
            or "\\" in self.storage_key
            or len(self.sha256) != 64
            or any(character not in "0123456789abcdef" for character in self.sha256)
            or not 0 < self.byte_size <= MAX_DOWNLOAD_BYTES
            or not self.mime_type
        ):
            raise PortalError("DOWNLOAD_REFERENCE_INVALID", "文件身份无效", 500)


class ArtifactReader(Protocol):
    def read(self, descriptor: ArtifactDescriptor) -> bytes:
        ...


class UnavailableArtifactReader:
    def read(self, descriptor: ArtifactDescriptor) -> bytes:
        raise PortalError("DOWNLOAD_UNAVAILABLE", "文件读取服务暂不可用", 503)


class ManagedObjectArtifactReader:
    """Read a committed daily XLSX from ARR's immutable object store."""

    def __init__(self, object_store: ManagedObjectStore) -> None:
        self._object_store = object_store

    def read(self, descriptor: ArtifactDescriptor) -> bytes:
        descriptor.validate()
        if descriptor.file_kind != "daily_xlsx":
            raise PortalError("DOWNLOAD_REFERENCE_INVALID", "文件身份无效", 500)
        try:
            stored = self._object_store.inspect_committed(
                descriptor.storage_key,
                descriptor.original_filename,
            )
            if (
                stored.role != "daily_report"
                or stored.sha256 != descriptor.sha256
                or stored.byte_size != descriptor.byte_size
                or stored.mime_type != descriptor.mime_type
            ):
                raise PortalError(
                    "DOWNLOAD_IDENTITY_MISMATCH", "文件完整性校验失败", 503
                )
            with tempfile.TemporaryDirectory(prefix="arr-download-") as temporary:
                destination = Path(temporary) / "daily.xlsx"
                self._object_store.materialize(
                    descriptor.storage_key,
                    destination,
                    MAX_DOWNLOAD_BYTES,
                )
                value = destination.read_bytes()
        except PortalError:
            raise
        except (IngestionError, OSError):
            raise PortalError(
                "DOWNLOAD_UNAVAILABLE", "文件读取服务暂不可用", 503
            ) from None
        if (
            len(value) != descriptor.byte_size
            or hashlib.sha256(value).hexdigest() != descriptor.sha256
        ):
            raise PortalError(
                "DOWNLOAD_IDENTITY_MISMATCH", "文件完整性校验失败", 503
            )
        return value


class RoutedArtifactReader:
    """Route OSS daily artifacts and controlled local report artifacts safely."""

    def __init__(
        self,
        *,
        daily_reader: Optional[ArtifactReader],
        local_reader: ArtifactReader,
    ) -> None:
        self._daily_reader = daily_reader
        self._local_reader = local_reader

    def read(self, descriptor: ArtifactDescriptor) -> bytes:
        descriptor.validate()
        if descriptor.file_kind == "daily_xlsx":
            if self._daily_reader is None:
                raise PortalError(
                    "DOWNLOAD_UNAVAILABLE", "文件读取服务暂不可用", 503
                )
            return self._daily_reader.read(descriptor)
        return self._local_reader.read(descriptor)


class ControlledProjectArtifactReader:
    """Development/local publication reader constrained to one project root."""

    def __init__(self, project_root: Path) -> None:
        self._project_root = project_root.resolve()

    def read(self, descriptor: ArtifactDescriptor) -> bytes:
        descriptor.validate()
        path = (self._project_root / descriptor.storage_key).resolve()
        try:
            path.relative_to(self._project_root)
            metadata = path.lstat()
        except (ValueError, OSError):
            raise PortalError("DOWNLOAD_UNAVAILABLE", "文件读取服务暂不可用", 503) from None
        if path.is_symlink() or not stat.S_ISREG(metadata.st_mode):
            raise PortalError("DOWNLOAD_REFERENCE_INVALID", "文件身份无效", 500)
        if metadata.st_size != descriptor.byte_size:
            raise PortalError("DOWNLOAD_IDENTITY_MISMATCH", "文件完整性校验失败", 503)
        digest = hashlib.sha256()
        chunks: list[bytes] = []
        descriptor_fd = -1
        total = 0
        try:
            descriptor_fd = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
            with os.fdopen(descriptor_fd, "rb") as source:
                descriptor_fd = -1
                while True:
                    chunk = source.read(1024 * 1024)
                    if not chunk:
                        break
                    total += len(chunk)
                    if total > MAX_DOWNLOAD_BYTES:
                        raise PortalError("DOWNLOAD_TOO_LARGE", "文件超过下载限制", 413)
                    digest.update(chunk)
                    chunks.append(chunk)
        except PortalError:
            raise
        except OSError:
            raise PortalError("DOWNLOAD_UNAVAILABLE", "文件读取服务暂不可用", 503) from None
        finally:
            if descriptor_fd >= 0:
                os.close(descriptor_fd)
        if total != descriptor.byte_size or digest.hexdigest() != descriptor.sha256:
            raise PortalError("DOWNLOAD_IDENTITY_MISMATCH", "文件完整性校验失败", 503)
        return b"".join(chunks)