redblack-python-sdk 0.1.0__py3-none-any.whl

redblack/downloads.py

@@ -0,0 +1,410 @@
+"""Credential-free, validated downloads for short-lived export URLs."""
+
+from __future__ import annotations
+
+import ipaddress
+import os
+import socket
+import tempfile
+from collections.abc import Callable, Iterable, Iterator, Mapping
+from contextlib import ExitStack, contextmanager, suppress
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, BinaryIO, cast
+from urllib.parse import urljoin, urlsplit, urlunsplit
+
+import httpx
+
+from .errors import RedBlackError
+
+_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
+_SAFE_DOWNLOAD_HEADERS = {
+    "Accept": "text/csv, application/octet-stream;q=0.9",
+    # Pinned URLs use the IP address as the HTTPX connection-pool origin.
+    # Prevent a redirect to another hostname on that IP from reusing TLS with
+    # the first hostname's SNI identity.
+    "Connection": "close",
+}
+
+
+class DownloadError(RedBlackError):
+    """Base class for export-storage download failures."""
+
+
+class DownloadSecurityError(DownloadError):
+    """The supplied storage URL violates the download security policy."""
+
+
+class DownloadExistsError(DownloadError):
+    """The destination already exists and overwrite was not enabled."""
+
+
+AddressResolver = Callable[[str, int], Iterable[str]]
+
+
+def _default_resolver(host: str, port: int) -> Iterable[str]:
+    addresses: list[str] = []
+    seen: set[str] = set()
+    for item in socket.getaddrinfo(host, port, type=socket.SOCK_STREAM):
+        address = cast(str, item[4][0])
+        if address not in seen:
+            seen.add(address)
+            addresses.append(address)
+    return addresses
+
+
+def _unsafe_address(address: str) -> bool:
+    parsed = ipaddress.ip_address(address)
+    return any(
+        (
+            parsed.is_private,
+            parsed.is_loopback,
+            parsed.is_link_local,
+            parsed.is_multicast,
+            parsed.is_reserved,
+            parsed.is_unspecified,
+        )
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class StorageUrlPolicy:
+    """Controls which unauthenticated storage URLs may be fetched."""
+
+    allowed_hosts: frozenset[str] | None = None
+    max_redirects: int = 3
+    resolve_hosts: bool = True
+
+    def __post_init__(self) -> None:
+        if self.max_redirects < 0:
+            raise ValueError("max_redirects must be greater than or equal to zero")
+        if self.allowed_hosts is not None:
+            normalized = frozenset(host.lower().rstrip(".") for host in self.allowed_hosts)
+            if not normalized or any(not host for host in normalized):
+                raise ValueError("allowed_hosts must contain valid host names")
+            object.__setattr__(self, "allowed_hosts", normalized)
+
+
+@dataclass(frozen=True, slots=True)
+class _ValidatedStorageTarget:
+    host: str
+    port: int
+    path: str
+    query: str
+    addresses: tuple[str, ...]
+
+    @property
+    def host_header(self) -> str:
+        host = f"[{self.host}]" if ":" in self.host else self.host
+        return f"{host}:{self.port}" if self.port != 443 else host
+
+    def pinned_url(self, address: str) -> str:
+        authority = f"[{address}]" if ":" in address else address
+        if self.port != 443:
+            authority = f"{authority}:{self.port}"
+        return urlunsplit(("https", authority, self.path, self.query, ""))
+
+
+def _canonical_host(host: str) -> str:
+    try:
+        return str(ipaddress.ip_address(host))
+    except ValueError:
+        pass
+    try:
+        return host.encode("idna").decode("ascii")
+    except UnicodeError as exc:
+        raise DownloadSecurityError("storage URL contains an invalid host") from exc
+
+
+def _validate_storage_target(
+    url: str,
+    *,
+    policy: StorageUrlPolicy,
+    resolver: AddressResolver | None,
+) -> _ValidatedStorageTarget:
+    if not isinstance(url, str) or not url:
+        raise DownloadSecurityError("storage URL is missing")
+    parsed = urlsplit(url)
+    if parsed.scheme.lower() != "https":
+        raise DownloadSecurityError("storage URL must use HTTPS")
+    if parsed.username is not None or parsed.password is not None:
+        raise DownloadSecurityError("storage URL must not contain user information")
+    if not parsed.hostname:
+        raise DownloadSecurityError("storage URL must include a host")
+    try:
+        parsed_port = parsed.port
+    except ValueError as exc:
+        raise DownloadSecurityError("storage URL contains an invalid port") from exc
+    port = 443 if parsed_port is None else parsed_port
+    if port == 0:
+        raise DownloadSecurityError("storage URL contains an invalid port")
+
+    source_host = parsed.hostname.lower().rstrip(".")
+    if not source_host:
+        raise DownloadSecurityError("storage URL must include a host")
+    host = _canonical_host(source_host)
+    if host == "localhost" or host.endswith(".localhost"):
+        raise DownloadSecurityError("local storage targets are not permitted")
+    if policy.allowed_hosts is not None and source_host not in policy.allowed_hosts:
+        raise DownloadSecurityError("storage host is not in the configured allowlist")
+
+    try:
+        literal = ipaddress.ip_address(host)
+    except ValueError:
+        literal = None
+    if literal is not None:
+        if _unsafe_address(str(literal)):
+            raise DownloadSecurityError("private or local storage targets are not permitted")
+        addresses = (str(literal),)
+    elif policy.resolve_hosts:
+        resolve = resolver or _default_resolver
+        try:
+            resolved = tuple(resolve(source_host, port))
+        except OSError as exc:
+            raise DownloadSecurityError("storage host could not be resolved") from exc
+        if not resolved:
+            raise DownloadSecurityError("storage host did not resolve to an address")
+
+        canonical_addresses: list[str] = []
+        seen_addresses: set[str] = set()
+        try:
+            for address in resolved:
+                canonical = str(ipaddress.ip_address(address))
+                if _unsafe_address(canonical):
+                    raise DownloadSecurityError(
+                        "storage host resolves to a private or local address"
+                    )
+                if canonical not in seen_addresses:
+                    seen_addresses.add(canonical)
+                    canonical_addresses.append(canonical)
+        except ValueError as exc:
+            raise DownloadSecurityError("storage host resolved to an invalid address") from exc
+        addresses = tuple(canonical_addresses)
+    else:
+        addresses = ()
+
+    return _ValidatedStorageTarget(
+        host=host,
+        port=port,
+        path=parsed.path,
+        query=parsed.query,
+        addresses=addresses,
+    )
+
+
+def validate_storage_url(
+    url: str,
+    *,
+    policy: StorageUrlPolicy | None = None,
+    resolver: AddressResolver | None = None,
+) -> str:
+    """Validate one storage URL and return it unchanged.
+
+    Every redirect target must be passed through this function again.  This
+    function intentionally accepts only HTTPS URLs without embedded user info.
+    """
+
+    _validate_storage_target(
+        url,
+        policy=policy or StorageUrlPolicy(),
+        resolver=resolver,
+    )
+    return url
+
+
+def _header(headers: Mapping[str, str], name: str) -> str | None:
+    target = name.casefold()
+    for key, value in headers.items():
+        if key.casefold() == target:
+            return value
+    return None
+
+
+@dataclass(frozen=True, slots=True)
+class DownloadResult:
+    path: Path
+    bytes_written: int
+    content_type: str | None
+
+
+class Downloader:
+    """Use a separate HTTP client so RedBlack credentials cannot leak."""
+
+    def __init__(
+        self,
+        *,
+        policy: StorageUrlPolicy | None = None,
+        client: Any | None = None,
+        resolver: AddressResolver | None = None,
+    ) -> None:
+        self._policy = policy or StorageUrlPolicy()
+        if client is not None:
+            self._validate_injected_client(client)
+        self._client = client
+        self._resolver = resolver
+
+    @staticmethod
+    def _validate_injected_client(client: object) -> None:
+        """Reject defaults that could be merged into a cross-origin request."""
+
+        sensitive = {
+            "accept-version",
+            "authorization",
+            "client-id",
+            "cookie",
+            "proxy-authorization",
+            "x-api-key",
+        }
+        headers = getattr(client, "headers", None)
+        if isinstance(headers, Mapping):
+            present = {str(name).casefold() for name in headers}
+            leaked = sensitive.intersection(present)
+            if leaked:
+                raise DownloadSecurityError(
+                    "injected download client has credential-bearing default headers"
+                )
+        cookies = getattr(client, "cookies", None)
+        if cookies is not None:
+            try:
+                has_cookies = len(cookies) > 0
+            except (TypeError, AttributeError):
+                has_cookies = bool(cookies)
+            if has_cookies:
+                raise DownloadSecurityError("injected download client must not contain cookies")
+        configured_auth = getattr(client, "_auth", None)
+        if configured_auth is not None:
+            raise DownloadSecurityError(
+                "injected download client must not have default authentication"
+            )
+
+    @contextmanager
+    def _pinned_response(
+        self,
+        client: Any,
+        target: _ValidatedStorageTarget,
+    ) -> Iterator[Any]:
+        if not target.addresses:
+            raise DownloadSecurityError(
+                "storage downloads require resolve_hosts=True for hostname pinning"
+            )
+
+        last_error: httpx.TransportError | None = None
+        for address in target.addresses:
+            headers = dict(_SAFE_DOWNLOAD_HEADERS)
+            headers["Host"] = target.host_header
+            with ExitStack() as stack:
+                try:
+                    response = stack.enter_context(
+                        client.stream(
+                            "GET",
+                            target.pinned_url(address),
+                            headers=headers,
+                            follow_redirects=False,
+                            # HTTPcore uses this extension for both TLS SNI and
+                            # certificate hostname verification while TCP is
+                            # connected to the vetted IP address.
+                            extensions={"sni_hostname": target.host},
+                        )
+                    )
+                except httpx.TransportError as exc:
+                    last_error = exc
+                    continue
+                yield response
+                return
+
+        raise DownloadError("storage download failed before receiving a response") from last_error
+
+    @contextmanager
+    def stream(
+        self, url: str, *, chunk_size: int = 64 * 1024
+    ) -> Iterator[tuple[Iterator[bytes], str | None]]:
+        """Yield a byte iterator and content type, closing all responses safely."""
+
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be greater than zero")
+        client = self._client or httpx.Client(
+            follow_redirects=False,
+            trust_env=False,
+            timeout=httpx.Timeout(30.0, connect=10.0),
+        )
+        owns_client = self._client is None
+        current = url
+        try:
+            for redirect_count in range(self._policy.max_redirects + 1):
+                target = _validate_storage_target(
+                    current,
+                    policy=self._policy,
+                    resolver=self._resolver,
+                )
+                with self._pinned_response(client, target) as response:
+                    if response.status_code in _REDIRECT_STATUSES:
+                        if redirect_count >= self._policy.max_redirects:
+                            raise DownloadError("storage redirect limit exceeded")
+                        location = _header(response.headers, "Location")
+                        if not location:
+                            raise DownloadError("storage redirect omitted Location")
+                        current = urljoin(current, location)
+                        continue
+                    if response.status_code < 200 or response.status_code >= 300:
+                        raise DownloadError(
+                            f"storage download failed with HTTP {response.status_code}"
+                        )
+                    yield response.iter_bytes(chunk_size), _header(response.headers, "Content-Type")
+                    return
+            raise DownloadError("storage redirect limit exceeded")
+        finally:
+            if owns_client:
+                client.close()
+
+    def download_to(
+        self,
+        url: str,
+        destination: str | os.PathLike[str],
+        *,
+        overwrite: bool = False,
+        chunk_size: int = 64 * 1024,
+    ) -> DownloadResult:
+        """Stream to a same-directory temporary file, then publish atomically."""
+
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be greater than zero")
+        target = Path(destination)
+        parent = target.parent
+        if not parent.exists():
+            raise FileNotFoundError(f"destination directory does not exist: {parent}")
+        if target.exists() and not overwrite:
+            raise DownloadExistsError(f"destination already exists: {target}")
+
+        descriptor, temporary_name = tempfile.mkstemp(
+            dir=parent, prefix=f".{target.name}.", suffix=".part"
+        )
+        temporary = Path(temporary_name)
+        bytes_written = 0
+        content_type: str | None = None
+        try:
+            with os.fdopen(descriptor, "wb") as output:
+                output_file: BinaryIO = output
+                with self.stream(url, chunk_size=chunk_size) as (chunks, content_type):
+                    for chunk in chunks:
+                        if not chunk:
+                            continue
+                        output_file.write(chunk)
+                        bytes_written += len(chunk)
+                output_file.flush()
+                os.fsync(output_file.fileno())
+
+            if overwrite:
+                os.replace(temporary, target)
+            else:
+                try:
+                    os.link(temporary, target)
+                except FileExistsError as exc:
+                    raise DownloadExistsError(f"destination already exists: {target}") from exc
+                temporary.unlink()
+            return DownloadResult(
+                path=target, bytes_written=bytes_written, content_type=content_type
+            )
+        except BaseException:
+            with suppress(OSError):
+                temporary.unlink(missing_ok=True)
+            raise