.venv/Lib/site-packages/pip/_internal/network/download.py

"""Download files with progress indicators."""

from __future__ import annotations

import email.message
import logging
import mimetypes
import os
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from http import HTTPStatus
from typing import BinaryIO

from pip._vendor.requests import PreparedRequest
from pip._vendor.requests.models import Response
from pip._vendor.urllib3 import HTTPResponse as URLlib3Response
from pip._vendor.urllib3._collections import HTTPHeaderDict
from pip._vendor.urllib3.exceptions import ReadTimeoutError

from pip._internal.cli.progress_bars import BarType, get_download_progress_renderer
from pip._internal.exceptions import IncompleteDownloadError, NetworkConnectionError
from pip._internal.models.index import PyPI
from pip._internal.models.link import Link
from pip._internal.network.cache import SafeFileCache, is_from_cache
from pip._internal.network.session import CacheControlAdapter, PipSession
from pip._internal.network.utils import HEADERS, raise_for_status, response_chunks
from pip._internal.utils.misc import format_size, redact_auth_from_url, splitext

logger = logging.getLogger(__name__)


def _get_http_response_size(resp: Response) -> int | None:
    try:
        return int(resp.headers["content-length"])
    except (ValueError, KeyError, TypeError):
        return None


def _get_http_response_etag_or_last_modified(resp: Response) -> str | None:
    """
    Return either the ETag or Last-Modified header (or None if neither exists).
    The return value can be used in an If-Range header.
    """
    return resp.headers.get("etag", resp.headers.get("last-modified"))


def _log_download(
    resp: Response,
    link: Link,
    progress_bar: BarType,
    total_length: int | None,
    range_start: int | None = 0,
) -> Iterable[bytes]:
    if link.netloc == PyPI.file_storage_domain:
        url = link.show_url
    else:
        url = link.url_without_fragment

    logged_url = redact_auth_from_url(url)

    if total_length:
        if range_start:
            logged_url = (
                f"{logged_url} ({format_size(range_start)}/{format_size(total_length)})"
            )
        else:
            logged_url = f"{logged_url} ({format_size(total_length)})"

    if is_from_cache(resp):
        logger.info("Using cached %s", logged_url)
    elif range_start:
        logger.info("Resuming download %s", logged_url)
    else:
        logger.info("Downloading %s", logged_url)

    if logger.getEffectiveLevel() > logging.INFO:
        show_progress = False
    elif is_from_cache(resp):
        show_progress = False
    elif not total_length:
        show_progress = True
    elif total_length > (512 * 1024):
        show_progress = True
    else:
        show_progress = False

    chunks = response_chunks(resp)

    if not show_progress:
        return chunks

    renderer = get_download_progress_renderer(
        bar_type=progress_bar, size=total_length, initial_progress=range_start
    )
    return renderer(chunks)


def sanitize_content_filename(filename: str) -> str:
    """
    Sanitize the "filename" value from a Content-Disposition header.
    """
    return os.path.basename(filename)


def parse_content_disposition(content_disposition: str, default_filename: str) -> str:
    """
    Parse the "filename" value from a Content-Disposition header, and
    return the default filename if the result is empty.
    """
    m = email.message.Message()
    m["content-type"] = content_disposition
    filename = m.get_param("filename")
    if filename:
        # We need to sanitize the filename to prevent directory traversal
        # in case the filename contains ".." path parts.
        filename = sanitize_content_filename(str(filename))
    return filename or default_filename


def _get_http_response_filename(resp: Response, link: Link) -> str:
    """Get an ideal filename from the given HTTP response, falling back to
    the link filename if not provided.
    """
    filename = link.filename  # fallback
    # Have a look at the Content-Disposition header for a better guess
    content_disposition = resp.headers.get("content-disposition")
    if content_disposition:
        filename = parse_content_disposition(content_disposition, filename)
    ext: str | None = splitext(filename)[1]
    if not ext:
        ext = mimetypes.guess_extension(resp.headers.get("content-type", ""))
        if ext:
            filename += ext
    if not ext and link.url != resp.url:
        ext = os.path.splitext(resp.url)[1]
        if ext:
            filename += ext
    return filename


@dataclass
class _FileDownload:
    """Stores the state of a single link download."""

    link: Link
    output_file: BinaryIO
    size: int | None
    bytes_received: int = 0
    reattempts: int = 0

    def is_incomplete(self) -> bool:
        return bool(self.size is not None and self.bytes_received < self.size)

    def write_chunk(self, data: bytes) -> None:
        self.bytes_received += len(data)
        self.output_file.write(data)

    def reset_file(self) -> None:
        """Delete any saved data and reset progress to zero."""
        self.output_file.seek(0)
        self.output_file.truncate()
        self.bytes_received = 0


class Downloader:
    def __init__(
        self,
        session: PipSession,
        progress_bar: BarType,
        resume_retries: int,
    ) -> None:
        assert (
            resume_retries >= 0
        ), "Number of max resume retries must be bigger or equal to zero"
        self._session = session
        self._progress_bar = progress_bar
        self._resume_retries = resume_retries

    def batch(
        self, links: Iterable[Link], location: str
    ) -> Iterable[tuple[Link, tuple[str, str]]]:
        """Convenience method to download multiple links."""
        for link in links:
            filepath, content_type = self(link, location)
            yield link, (filepath, content_type)

    def __call__(self, link: Link, location: str) -> tuple[str, str]:
        """Download a link and save it under location."""
        resp = self._http_get(link)
        download_size = _get_http_response_size(resp)

        filepath = os.path.join(location, _get_http_response_filename(resp, link))
        with open(filepath, "wb") as content_file:
            download = _FileDownload(link, content_file, download_size)
            self._process_response(download, resp)
            if download.is_incomplete():
                self._attempt_resumes_or_redownloads(download, resp)

        content_type = resp.headers.get("Content-Type", "")
        return filepath, content_type

    def _process_response(self, download: _FileDownload, resp: Response) -> None:
        """Download and save chunks from a response."""
        chunks = _log_download(
            resp,
            download.link,
            self._progress_bar,
            download.size,
            range_start=download.bytes_received,
        )
        try:
            for chunk in chunks:
                download.write_chunk(chunk)
        except ReadTimeoutError as e:
            # If the download size is not known, then give up downloading the file.
            if download.size is None:
                raise e

            logger.warning("Connection timed out while downloading.")

    def _attempt_resumes_or_redownloads(
        self, download: _FileDownload, first_resp: Response
    ) -> None:
        """Attempt to resume/restart the download if connection was dropped."""

        while download.reattempts < self._resume_retries and download.is_incomplete():
            assert download.size is not None
            download.reattempts += 1
            logger.warning(
                "Attempting to resume incomplete download (%s/%s, attempt %d)",
                format_size(download.bytes_received),
                format_size(download.size),
                download.reattempts,
            )

            try:
                resume_resp = self._http_get_resume(download, should_match=first_resp)
                # Fallback: if the server responded with 200 (i.e., the file has
                # since been modified or range requests are unsupported) or any
                # other unexpected status, restart the download from the beginning.
                must_restart = resume_resp.status_code != HTTPStatus.PARTIAL_CONTENT
                if must_restart:
                    download.reset_file()
                    download.size = _get_http_response_size(resume_resp)
                    first_resp = resume_resp

                self._process_response(download, resume_resp)
            except (ConnectionError, ReadTimeoutError, OSError):
                continue

        # No more resume attempts. Raise an error if the download is still incomplete.
        if download.is_incomplete():
            os.remove(download.output_file.name)
            raise IncompleteDownloadError(download)

        # If we successfully completed the download via resume, manually cache it
        # as a complete response to enable future caching
        if download.reattempts > 0:
            self._cache_resumed_download(download, first_resp)

    def _cache_resumed_download(
        self, download: _FileDownload, original_response: Response
    ) -> None:
        """
        Manually cache a file that was successfully downloaded via resume retries.

        cachecontrol doesn't cache 206 (Partial Content) responses, since they
        are not complete files. This method manually adds the final file to the
        cache as though it was downloaded in a single request, so that future
        requests can use the cache.
        """
        url = download.link.url_without_fragment
        adapter = self._session.get_adapter(url)

        # Check if the adapter is the CacheControlAdapter (i.e. caching is enabled)
        if not isinstance(adapter, CacheControlAdapter):
            logger.debug(
                "Skipping resume download caching: no cache controller for %s", url
            )
            return

        # Check SafeFileCache is being used
        assert isinstance(
            adapter.cache, SafeFileCache
        ), "separate body cache not in use!"

        synthetic_request = PreparedRequest()
        synthetic_request.prepare(method="GET", url=url, headers={})

        synthetic_response_headers = HTTPHeaderDict()
        for key, value in original_response.headers.items():
            if key.lower() not in ["content-range", "content-length"]:
                synthetic_response_headers[key] = value
        synthetic_response_headers["content-length"] = str(download.size)

        synthetic_response = URLlib3Response(
            body="",
            headers=synthetic_response_headers,
            status=200,
            preload_content=False,
        )

        # Save metadata and then stream the file contents to cache.
        cache_url = adapter.controller.cache_url(url)
        metadata_blob = adapter.controller.serializer.dumps(
            synthetic_request, synthetic_response, b""
        )
        adapter.cache.set(cache_url, metadata_blob)
        download.output_file.flush()
        with open(download.output_file.name, "rb") as f:
            adapter.cache.set_body_from_io(cache_url, f)

        logger.debug(
            "Cached resumed download as complete response for future use: %s", url
        )

    def _http_get_resume(
        self, download: _FileDownload, should_match: Response
    ) -> Response:
        """Issue a HTTP range request to resume the download."""
        # To better understand the download resumption logic, see the mdn web docs:
        # https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/Range_requests
        headers = HEADERS.copy()
        headers["Range"] = f"bytes={download.bytes_received}-"
        # If possible, use a conditional range request to avoid corrupted
        # downloads caused by the remote file changing in-between.
        if identifier := _get_http_response_etag_or_last_modified(should_match):
            headers["If-Range"] = identifier
        return self._http_get(download.link, headers)

    def _http_get(self, link: Link, headers: Mapping[str, str] = HEADERS) -> Response:
        target_url = link.url_without_fragment
        try:
            resp = self._session.get(target_url, headers=headers, stream=True)
            raise_for_status(resp)
        except NetworkConnectionError as e:
            assert e.response is not None
            logger.critical(
                "HTTP error %s while getting %s", e.response.status_code, link
            )
            raise
        return resp