test_docker_pull.py

"""Pytest coverage for the Docker image downloader."""

import gzip
import hashlib
import io
import json
import os
import subprocess
import sys
import tarfile
import threading
from email.message import Message
from http.client import HTTPMessage
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import patch
from urllib.error import HTTPError
from urllib.request import Request, build_opener

import pytest

import docker_pull
from docker_pull import (
    Config,
    DockerImagePuller,
    OriginSafeRedirectHandler,
    ProgressReporter,
    ProxyManager,
    RegistryError,
    format_bytes,
    normalize_architecture,
    parse_image_spec,
    validate_digest,
)


class Response(io.BytesIO):
    """Small context-managed HTTP response used by offline tests."""

    def __init__(self, body=b"", url="https://registry.example/resource", headers=None):
        super().__init__(body)
        self.code = 200
        self.headers = headers or {}
        self.url = url

    def geturl(self):
        return self.url


@pytest.fixture(autouse=True)
def clear_proxy_environment(monkeypatch):
    """Keep tests independent of the host's corporate proxy settings."""
    for name in (
        "HTTP_PROXY",
        "http_proxy",
        "HTTPS_PROXY",
        "https_proxy",
        "NO_PROXY",
        "no_proxy",
    ):
        monkeypatch.delenv(name, raising=False)


def content_digest(data, algorithm="sha256"):
    return f"{algorithm}:{hashlib.new(algorithm, data).hexdigest()}"


def offline_puller(manifest, config_data, downloaded=None):
    """Return a puller whose registry responses are supplied in memory."""
    downloaded = downloaded if downloaded is not None else []
    config_digest = manifest["config"]["digest"]
    puller = DockerImagePuller()
    puller.get_auth_header = lambda image_name, challenge=None: None
    puller.get_manifest = lambda *_args, **_kwargs: manifest

    def download_blob(
        image_name,
        digest,
        authorization,
        retry_with_new_auth=True,
        expected_size=None,
        destination=None,
    ):
        downloaded.append(digest)
        return config_data if digest == config_digest else None

    puller.download_blob = download_blob
    return puller


def test_config_validation(monkeypatch):
    config = Config(
        timeout_config={
            "request_timeout": 30,
            "download_timeout": 300,
            "chunk_timeout": 60,
        }
    )
    assert config.request_timeout == 30
    assert config.download_timeout == 300
    assert config.chunk_timeout == 60

    for timeout in (
        {"request_timeout": -1},
        {"request_timeout": "30"},
        {"download_timeout": True},
        {"download_timeout": 0},
        {"chunk_timeout": 1.5},
        {"chunk_timeout": True},
    ):
        with pytest.raises(ValueError):
            Config(timeout_config=timeout)

    caller_proxy_config = {"http_proxy": "http://proxy.example:8080"}
    Config(proxy_config=caller_proxy_config)
    assert caller_proxy_config == {"http_proxy": "http://proxy.example:8080"}

    monkeypatch.setenv("HTTP_PROXY", "http://test-proxy:8080")
    assert Config().proxy_config["http_proxy"] == "http://test-proxy:8080"

    with pytest.raises(ValueError, match="username and password"):
        DockerImagePuller(username="user")

    for url in (
        "http://proxy.example:invalid",
        "http://proxy.example:0",
        "http://proxy.example:",
    ):
        with pytest.raises(ValueError, match="Invalid http_proxy URL"):
            Config(proxy_config={"http_proxy": url})


def test_http_error_body_is_bounded_and_escaped():
    error = HTTPError(
        "https://registry.example/v2/team/app/manifests/latest",
        500,
        "Server Error",
        Message(),
        io.BytesIO(b"line1\n\x1b[31m" + b"x" * 5000),
    )
    with patch.object(docker_pull.logger, "error") as logged:
        docker_pull.log_http_error_body(error)

    template, detail = logged.call_args.args
    rendered = template % detail
    assert error.fp.tell() == 4096
    assert len(detail) == 500
    assert "\n" not in rendered
    assert "\x1b" not in rendered


def test_proxy_manager_security():
    manager = ProxyManager(Config())
    assert manager.sanitize_proxy_url("http://user:pass@proxy.com:8080") == (
        "http://proxy.com:8080"
    )
    assert manager.sanitize_proxy_url("https://user:pass@proxy.com") == (
        "https://proxy.com"
    )
    assert manager.sanitize_proxy_url("http://user:pass@[::1]:8080") == (
        "http://[::1]:8080"
    )
    assert manager.sanitize_proxy_url("http://proxy.com:8080") == (
        "http://proxy.com:8080"
    )
    assert manager.sanitize_proxy_url("") == ""
    assert manager.sanitize_proxy_url(None) is None
    assert "secret" not in manager._mask_credentials_fallback(
        "Connection to http://user:secret@proxy.com failed"
    )
    assert manager.sanitize_debug_output(
        {
            "Authorization": "Bearer secret",
            "Proxy-Authorization": "Basic secret",
            "Accept": "application/json",
        }
    ) == {
        "Authorization": "***",
        "Proxy-Authorization": "***",
        "Accept": "application/json",
    }
    assert manager._add_proxy_auth("http://proxy.example", "user:secret") == (
        "http://user:secret@proxy.example"
    )
    authenticated_proxy = manager._add_proxy_auth(
        "http://proxy.example:8080", "corp user:S3cr3t/Pa55?#"
    )
    assert authenticated_proxy == (
        "http://corp%20user:S3cr3t%2FPa55%3F%23@proxy.example:8080"
    )
    assert manager.sanitize_proxy_url(authenticated_proxy) == (
        "http://proxy.example:8080"
    )

    bypass_manager = ProxyManager(Config(proxy_config={"no_proxy": ".corp.com"}))
    assert bypass_manager.should_bypass_proxy("corp.com")
    assert bypass_manager.should_bypass_proxy("registry.corp.com")
    assert not bypass_manager.should_bypass_proxy("notcorp.com")

    port_bypass_manager = ProxyManager(
        Config(proxy_config={"no_proxy": "localhost:5000,[::1]:5000"})
    )
    assert port_bypass_manager.should_bypass_proxy("localhost:5000")
    assert not port_bypass_manager.should_bypass_proxy("localhost:6000")
    assert port_bypass_manager.should_bypass_proxy("[::1]:5000")

    with patch.object(docker_pull.logger, "warning") as warning:
        ProxyManager(Config(proxy_config={"insecure": True}))
    warning.assert_called_once_with("SSL certificate verification is disabled")


def test_progress_reporter():
    reporter = ProgressReporter(total_size=1000, description="Test", show_speed=False)
    assert reporter.downloaded == 0
    assert reporter.total_size == 1000
    assert format_bytes(512) == "512 B"
    assert format_bytes(1024) == "1.0 KB"
    assert format_bytes(1048576) == "1.0 MB"
    assert format_bytes(1073741824) == "1.0 GB"
    assert reporter._format_duration(30) == "30s"
    assert reporter._format_duration(90) == "1m30s"
    assert reporter._format_duration(3661) == "1h01m"

    reporter.downloaded = 512
    with (
        patch.object(docker_pull.logger, "isEnabledFor", return_value=False),
        patch("builtins.print") as output,
    ):
        reporter._display_progress()
        reporter.finish()
    output.assert_not_called()


def test_default_archive_mode_honors_posix_umask():
    with (
        patch.object(docker_pull.os, "name", "posix"),
        patch.object(docker_pull.os, "umask", side_effect=[0o027, 0]) as umask,
        patch.object(docker_pull.os, "chmod") as chmod,
    ):
        docker_pull.apply_default_file_mode("image.tar")
    assert umask.call_args_list[1].args == (0o027,)
    chmod.assert_called_once_with("image.tar", 0o640)


@pytest.mark.parametrize(
    ("image_spec", "expected"),
    [
        (
            "ubuntu",
            (
                "https://registry-1.docker.io",
                "library/ubuntu",
                "latest",
                "ubuntu",
                "latest",
            ),
        ),
        (
            "icr.io/codeengine/helloworld:latest",
            (
                "https://icr.io",
                "codeengine/helloworld",
                "latest",
                "icr.io/codeengine/helloworld",
                "latest",
            ),
        ),
        (
            "localhost:5000/myorg/myapp",
            (
                "https://localhost:5000",
                "myorg/myapp",
                "latest",
                "localhost:5000/myorg/myapp",
                "latest",
            ),
        ),
    ],
)
def test_image_spec_parsing(image_spec, expected):
    assert parse_image_spec(image_spec) == expected


@pytest.mark.parametrize(
    ("image_spec", "registry"),
    [
        ("ghcr.io/nextcloud-releases/all-in-one:latest", "https://ghcr.io"),
        ("registry.gitlab.com/group/project/image:tag", "https://registry.gitlab.com"),
        ("quay.io/prometheus/busybox:latest", "https://quay.io"),
        ("public.ecr.aws/amazonlinux/amazonlinux:latest", "https://public.ecr.aws"),
        ("us-docker.pkg.dev/project/repository/image:tag", "https://us-docker.pkg.dev"),
        ("example.azurecr.io/repository/image:tag", "https://example.azurecr.io"),
        ("icr.io/namespace/image:tag", "https://icr.io"),
        ("us-ashburn-1.ocir.io/tenancy/image:tag", "https://us-ashburn-1.ocir.io"),
        (
            "registry.digitalocean.com/example/image:tag",
            "https://registry.digitalocean.com",
        ),
        ("harbor.example.com/project/image:tag", "https://harbor.example.com"),
        ("localhost:5000/image:tag", "https://localhost:5000"),
        ("[::1]:5000/image:tag", "https://[::1]:5000"),
        ("nexus.example.com:5001/image:tag", "https://nexus.example.com:5001"),
        (
            "artifactory.example.com/repository/image:tag",
            "https://artifactory.example.com",
        ),
        ("quay.example.com/project/image:tag", "https://quay.example.com"),
    ],
)
def test_registry_parsing(image_spec, registry):
    assert parse_image_spec(image_spec)[0] == registry


def test_digest_and_image_spec_validation():
    digest = "sha256:" + "a" * 64
    assert parse_image_spec(f"ubuntu@{digest}") == (
        "https://registry-1.docker.io",
        "library/ubuntu",
        digest,
        "ubuntu",
        None,
    )
    assert parse_image_spec(f"ghcr.io/org/app:stable@{digest}") == (
        "https://ghcr.io",
        "org/app",
        digest,
        "ghcr.io/org/app",
        "stable",
    )
    assert parse_image_spec("localhost:5000/app", plain_http=True)[0] == (
        "http://localhost:5000"
    )
    assert parse_image_spec("ubuntu@sha256:" + "A" * 64)[2] == digest

    for image_spec in (
        "ubuntu@sha256:short",
        "ubuntu:",
        "ghcr.io/",
        "ubuntu:..\\evil",
        "registry.example/te am/app:v1",
        "registry.example/Team/app:v1",
        "registry.example/team/../../app:v1",
        "registry.example:invalid/team/app:v1",
        "host:/app:v1",
        "registry.example\n.evil/team/app:v1",
        "[not-ipv6]:5000/team/app:v1",
    ):
        with pytest.raises((RegistryError, ValueError)):
            parse_image_spec(image_spec)
    with pytest.raises(ValueError, match=r"\[@digest\]"):
        parse_image_spec("user:pass@registry.example/app:latest")
    with pytest.raises(ValueError, match=r"\[@digest\]"):
        parse_image_spec("user:pass@ubuntu")
    with pytest.raises(ValueError, match="must be lowercase"):
        parse_image_spec("ghcr.io/NextCloud/aio:latest")
    for digest in (
        "sha256:/tmp/owned",
        "md5:" + "0" * 32,
        "SHA256:" + "0" * 64,
        123,
    ):
        with pytest.raises(RegistryError):
            validate_digest(digest)


def test_registry_authentication():
    requested = []
    challenge = (
        'Bearer realm="https://login.example.com/token",'
        'service="registry.example.com",scope="repository:user/image:pull"'
    )

    def make_request(url, headers=None):
        requested.append((url, headers))
        if url.endswith("/v2/"):
            response_headers = Message()
            response_headers["WWW-Authenticate"] = challenge
            raise HTTPError(url, 401, "Unauthorized", response_headers, None)
        return Response(b'{"access_token":"registry-token"}')

    puller = DockerImagePuller()
    puller.make_request = make_request
    assert puller.get_auth_header("team/app") == "Bearer registry-token"
    assert "service=registry.example.com" in requested[1][0]
    assert "scope=repository%3Ateam%2Fapp%3Apull" in requested[1][0]

    basic_puller = DockerImagePuller(username="user", password="secret")
    assert basic_puller.get_auth_header("team/app", 'Basic realm="Registry"') == (
        "Basic dXNlcjpzZWNyZXQ="
    )

    combined_puller = DockerImagePuller(username="user", password="secret")
    combined_puller.make_request = lambda url, headers=None: Response(
        b'{"token":"registry-token"}'
    )
    combined = (
        'Basic realm="Registry", Bearer realm="https://login.example/token",'
        "service=registry.example"
    )
    assert combined_puller.get_auth_header("team/app", combined) == (
        "Bearer registry-token"
    )

    with pytest.raises(RegistryError, match="insecure authentication realm"):
        basic_puller.get_auth_header(
            "team/app",
            'Bearer realm="http://login.example/token",service="registry.example"',
        )

    malformed_puller = DockerImagePuller()
    malformed_puller.make_request = lambda *_args, **_kwargs: Response(b"[]")
    with pytest.raises(RegistryError, match="token response is not a JSON object"):
        malformed_puller.get_auth_header(
            "team/app", 'Bearer realm="https://login.example/token"'
        )

    malformed_puller.make_request = lambda *_args, **_kwargs: Response(
        b'{"token":"line1\\n line2"}'
    )
    with pytest.raises(RegistryError, match="multi-line token"):
        malformed_puller.get_auth_header(
            "team/app", 'Bearer realm="https://login.example/token"'
        )

    with pytest.raises(RegistryError, match="Malformed authentication challenge"):
        malformed_puller.get_auth_header("team/app", "Bearer realm=")


def test_auth_challenge_parsing_respects_quotes_and_multiple_fields():
    parsed = docker_pull.parse_auth_challenges(
        'Bearer realm="https://login.example/token,Basic fake",service="registry"'
    )
    assert parsed == [
        (
            "bearer",
            {
                "realm": "https://login.example/token,Basic fake",
                "service": "registry",
            },
        )
    ]

    headers = Message()
    headers["WWW-Authenticate"] = 'Basic realm="Registry"'
    headers["WWW-Authenticate"] = (
        'Bearer realm="https://login.example/token",service="registry"'
    )
    challenges = docker_pull.get_auth_challenges(headers)
    assert len(challenges) == 2
    assert docker_pull.parse_auth_challenges(challenges)[0][0] == "bearer"


def credential_headers(request):
    return {
        key.lower(): value
        for store in (request.headers, request.unredirected_hdrs)
        for key, value in store.items()
        if key.lower() in ("authorization", "proxy-authorization")
    }


@pytest.mark.parametrize("code", [301, 302, 303, 307, 308])
def test_redirect_security(code):
    handler = OriginSafeRedirectHandler()
    response_headers = HTTPMessage()
    request = Request(
        "https://registry.example/v2/blob",
        headers={
            "Authorization": "Bearer secret",
            "Proxy-Authorization": "Basic proxy-secret",
        },
    )

    same_origin = handler.redirect_request(
        request,
        io.BytesIO(),
        code,
        "Redirect",
        response_headers,
        "https://registry.example/storage/blob",
    )
    assert credential_headers(same_origin) == {"authorization": "Bearer secret"}

    cross_origin = handler.redirect_request(
        request,
        io.BytesIO(),
        code,
        "Redirect",
        response_headers,
        "https://cdn.example/blob",
    )
    assert cross_origin is not None
    assert credential_headers(cross_origin) == {}

    second_hop = handler.redirect_request(
        cross_origin,
        io.BytesIO(),
        code,
        "Redirect",
        response_headers,
        "https://cdn.example/blob-2",
    )
    assert credential_headers(second_hop) == {}


def test_unsafe_redirects_are_rejected():
    handler = OriginSafeRedirectHandler()
    headers = HTTPMessage()
    with pytest.raises(HTTPError, match="unsupported scheme"):
        handler.redirect_request(
            Request("http://registry.example/v2/blob"),
            io.BytesIO(),
            302,
            "Found",
            headers,
            "ftp://cdn.example/blob",
        )
    with pytest.raises(HTTPError, match="HTTPS downgrade"):
        handler.redirect_request(
            Request("https://registry.example/v2/blob"),
            io.BytesIO(),
            302,
            "Found",
            headers,
            "http://registry.example/blob",
        )
    with pytest.raises(HTTPError, match="invalid URL"):
        handler.redirect_request(
            Request("https://registry.example/v2/blob"),
            io.BytesIO(),
            302,
            "Found",
            headers,
            "https://cdn.example:70000/blob",
        )


def test_redirect_security_through_real_opener():
    seen = []

    class CdnHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            seen.append(
                (
                    "cdn",
                    self.headers.get("Authorization"),
                    self.headers.get("Proxy-Authorization"),
                )
            )
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"ok")

        def log_message(self, format, *args):
            del format, args

    class RegistryHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            seen.append(
                (
                    self.path,
                    self.headers.get("Authorization"),
                    self.headers.get("Proxy-Authorization"),
                )
            )
            if self.path == "/start":
                location, code = "/same-origin", 302
            else:
                host, port = cdn.server_address[0], cdn.server_address[1]
                location, code = f"http://{host}:{port}/final", 307
            self.send_response(code)
            self.send_header("Location", location)
            self.end_headers()

        def log_message(self, format, *args):
            del format, args

    with (
        ThreadingHTTPServer(("127.0.0.1", 0), CdnHandler) as cdn,
        ThreadingHTTPServer(("127.0.0.1", 0), RegistryHandler) as registry,
    ):
        threads = [
            threading.Thread(target=server.serve_forever, daemon=True)
            for server in (cdn, registry)
        ]
        for thread in threads:
            thread.start()
        try:
            host, port = registry.server_address[0], registry.server_address[1]
            request = Request(
                f"http://{host}:{port}/start",
                headers={
                    "Authorization": "Bearer secret",
                    "Proxy-Authorization": "Basic proxy-secret",
                },
            )
            with build_opener(OriginSafeRedirectHandler()).open(
                request, timeout=5
            ) as response:
                assert response.read() == b"ok"
        finally:
            for server in (cdn, registry):
                server.shutdown()
            for thread in threads:
                thread.join(timeout=5)

    assert seen == [
        ("/start", "Bearer secret", "Basic proxy-secret"),
        ("/same-origin", "Bearer secret", None),
        ("cdn", None, None),
    ]


def test_blob_validation_and_cleanup(tmp_path):
    puller = DockerImagePuller()
    data = b"verified blob"
    digest = content_digest(data)
    assert puller._stream_download(io.BytesIO(data), digest, len(data)) == data
    assert puller._stream_download(io.BytesIO(data), "sha256:" + "0" * 64) is None
    assert puller._stream_download(io.BytesIO(data), digest, len(data) + 1) is None

    oversized = io.BytesIO(b"x" * (65536 * 2))
    assert puller._stream_download(oversized, content_digest(b"x"), 1) is None
    assert oversized.tell() == 65536

    destination = tmp_path / "blob"
    assert puller._stream_download(
        io.BytesIO(data), digest, len(data), str(destination)
    ) == str(destination)
    assert destination.read_bytes() == data
    assert not (tmp_path / "blob.part").exists()

    threaded_result = []
    worker = threading.Thread(
        target=lambda: threaded_result.append(
            puller._stream_download(io.BytesIO(data), digest, len(data))
        )
    )
    worker.start()
    worker.join()
    assert threaded_result == [data]

    with (
        patch.object(docker_pull.signal, "SIGALRM", 14, create=True),
        patch.object(docker_pull.signal, "alarm", create=True),
        patch.object(docker_pull.signal, "signal", return_value=None) as set_handler,
    ):
        assert puller._stream_download(io.BytesIO(data), digest, len(data)) == data
    assert set_handler.call_count == 1

    large_data = b"x" * (1024 * 1024 + 1)
    with (
        patch.object(docker_pull.logger, "isEnabledFor", return_value=False),
        patch("builtins.print") as output,
    ):
        assert (
            puller._stream_download(io.BytesIO(large_data), content_digest(large_data))
            == large_data
        )
    output.assert_not_called()

    with (
        patch.object(docker_pull.logger, "isEnabledFor", return_value=True),
        patch.object(docker_pull.time, "monotonic", return_value=1.0),
        patch("builtins.print") as output,
    ):
        assert (
            puller._stream_download(io.BytesIO(large_data), content_digest(large_data))
            == large_data
        )
    assert output.call_count == 2


def test_blob_401_refreshes_authentication():
    data = b"verified blob"
    digest = content_digest(data)
    response_headers = Message()
    response_headers["WWW-Authenticate"] = 'Bearer realm="https://login.example/token"'
    requests = []
    puller = DockerImagePuller()

    def open_request(request, *_args, **_kwargs):
        requests.append(request.get_header("Authorization"))
        if len(requests) == 1:
            raise HTTPError(
                request.full_url, 401, "Unauthorized", response_headers, None
            )
        return Response(data, headers={"Content-Length": str(len(data))})

    puller.proxy_manager.open = open_request
    puller.get_auth_header = lambda *_args, **_kwargs: "Bearer fresh"
    assert puller.download_blob("team/app", digest, "Bearer stale") == data
    assert requests == ["Bearer stale", "Bearer fresh"]


def test_redirected_blob_url_is_sanitized():
    data = b"verified blob"
    digest = content_digest(data)
    response = Response(
        data,
        url="https://cdn.example/blob?X-Amz-Credential=secret",
        headers={"Content-Length": str(len(data))},
    )
    puller = DockerImagePuller()
    puller.proxy_manager.open = lambda *_args, **_kwargs: response

    with patch.object(docker_pull.logger, "debug") as debug:
        assert puller.download_blob("team/app", digest, None) == data
    assert "X-Amz-Credential" not in " ".join(
        str(call) for call in debug.call_args_list
    )


def test_blob_url_uses_normalized_digest():
    data = b"verified blob"
    digest = content_digest(data)
    requested = []
    puller = DockerImagePuller()

    def open_request(request, *_args, **_kwargs):
        requested.append(request.full_url)
        return Response(data, headers={"Content-Length": str(len(data))})

    puller.proxy_manager.open = open_request
    assert (
        puller.download_blob(
            "team/app", digest.upper().replace("SHA256", "sha256"), None
        )
        == data
    )
    assert requested == [f"https://registry-1.docker.io/v2/team/app/blobs/{digest}"]


@pytest.mark.parametrize("content_length", ["garbage", "-1"])
def test_invalid_blob_content_length_is_ignored(content_length):
    data = b"verified blob"
    digest = content_digest(data)
    puller = DockerImagePuller()
    puller.proxy_manager.open = lambda *_args, **_kwargs: Response(
        data, headers={"Content-Length": content_length}
    )
    assert puller.download_blob("team/app", digest, None) == data


@pytest.mark.parametrize(
    ("architecture", "variant", "expected"),
    [
        ("aarch64", None, "arm64"),
        ("x86_64", None, "amd64"),
        ("x86-64", None, "amd64"),
        ("i386", None, "386"),
        ("arm", "v8", "arm64"),
        ("arm", "v7", "arm"),
    ],
)
def test_architecture_normalization(architecture, variant, expected):
    assert normalize_architecture(architecture, variant) == expected


@pytest.mark.parametrize(
    ("architecture", "variant"), [(123, None), (["amd64"], None), ("arm", 8)]
)
def test_malformed_config_architecture_is_a_registry_error(architecture, variant):
    config_data = json.dumps(
        {"architecture": architecture, "variant": variant, "os": "linux"}
    ).encode()
    config_digest = content_digest(config_data)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [{"digest": "sha256:" + "1" * 64, "size": 1}],
    }
    with pytest.raises(RegistryError, match="platform is invalid"):
        offline_puller(manifest, config_data).pull_image(
            "registry.example/team/app:latest"
        )


@pytest.mark.parametrize(
    ("index_architecture", "variant", "requested_architecture", "requested_variant"),
    [
        ("aarch64", None, "arm64", None),
        ("arm", "v8", "arm64", None),
        ("arm64", None, "arm64", "v8"),
    ],
)
def test_index_architecture_alias_selection(
    index_architecture, variant, requested_architecture, requested_variant
):
    manifest = {"schemaVersion": 2, "config": {}, "layers": []}
    manifest_body = json.dumps(manifest).encode()
    digest = content_digest(manifest_body)
    index = {
        "schemaVersion": 2,
        "manifests": [
            {
                "digest": digest,
                "size": len(manifest_body),
                "platform": {
                    "os": "linux",
                    "architecture": index_architecture,
                    "variant": variant,
                },
            }
        ],
    }
    responses = iter([Response(json.dumps(index).encode()), Response(manifest_body)])
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: next(responses)
    assert (
        puller.get_manifest(
            "team/app",
            "latest",
            None,
            requested_architecture,
            "linux",
            requested_variant,
        )
        == manifest
    )


def test_index_variant_and_os_selection():
    selected_manifest = {
        "schemaVersion": 2,
        "config": {},
        "layers": [],
        "selected": "v7",
    }
    selected_body = json.dumps(selected_manifest).encode()
    selected_digest = content_digest(selected_body)
    other_body = json.dumps({**selected_manifest, "selected": "v6"}).encode()
    index = {
        "schemaVersion": 2,
        "manifests": [
            {
                "digest": content_digest(other_body),
                "size": len(other_body),
                "platform": {"os": "linux", "architecture": "arm", "variant": "v6"},
            },
            {
                "digest": selected_digest,
                "size": len(selected_body),
                "platform": {"os": "Linux", "architecture": "arm", "variant": "V7"},
            },
        ],
    }
    requested = []
    responses = iter([Response(json.dumps(index).encode()), Response(selected_body)])
    puller = DockerImagePuller()

    def make_request(url, *_args, **_kwargs):
        requested.append(url)
        return next(responses)

    puller.make_request = make_request
    assert (
        puller.get_manifest("team/app", "latest", None, "arm", "linux", "v7")
        == selected_manifest
    )
    assert requested[-1].endswith(selected_digest)


def test_missing_platform_is_rejected(caplog):
    body = json.dumps(
        {
            "manifests": [
                {
                    "digest": "sha256:" + "a" * 64,
                    "platform": {"os": "linux", "architecture": "arm64"},
                }
            ]
        }
    ).encode()
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: Response(body)
    with pytest.raises(RegistryError, match="No manifest found"):
        puller.get_manifest("team/app", "latest", None, "amd64", "linux")

    malformed_body = json.dumps(
        {"manifests": [{"platform": {"os": "linux", "architecture": 123}}]}
    ).encode()
    puller.make_request = lambda *_args, **_kwargs: Response(malformed_body)
    with pytest.raises(RegistryError, match="1 malformed entry omitted"):
        puller.get_manifest("team/app", "latest", None, "amd64", "linux")

    injected = "linux\n\x1b[31mINJECTED\x1b[0m"
    hostile_body = json.dumps(
        {
            "manifests": [
                {
                    "digest": "sha256:" + "a" * 64,
                    "platform": {"os": injected, "architecture": "arm64"},
                }
            ]
        }
    ).encode()
    puller.make_request = lambda *_args, **_kwargs: Response(hostile_body)
    with pytest.raises(RegistryError) as error:
        puller.get_manifest("team/app", "latest", None, "amd64", "linux")
    assert "\n" not in str(error.value) and "\x1b" not in str(error.value)
    assert all(
        "\n" not in record.getMessage() and "\x1b" not in record.getMessage()
        for record in caplog.records
    )


def test_selected_manifest_requires_a_digest():
    body = json.dumps(
        {"manifests": [{"platform": {"os": "linux", "architecture": "amd64"}}]}
    ).encode()
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: Response(body)
    with pytest.raises(RegistryError, match="missing a digest"):
        puller.get_manifest("team/app", "latest", None)


@pytest.mark.parametrize(
    "body",
    [
        b"[]",
        b'{"manifests":"invalid"}',
        b'{"manifests":["invalid"]}',
        b'{"manifests":[{"platform":"linux"}]}',
        b'{"schemaVersion":2,"mediaType":123,"config":{}}',
    ],
)
def test_malformed_manifest_json_is_a_registry_error(body):
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: Response(body)
    with pytest.raises(RegistryError):
        puller.get_manifest("team/app", "latest", None)


def test_manifest_integrity():
    manifest = {"schemaVersion": 2, "config": {}, "layers": []}
    manifest_body = json.dumps(manifest).encode()
    manifest_digest = content_digest(manifest_body)
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: Response(manifest_body)
    assert puller.get_manifest("team/app", manifest_digest, None) == manifest
    with pytest.raises(RegistryError, match="digest mismatch"):
        puller.get_manifest("team/app", "sha256:" + "0" * 64, None)

    requested = []
    uppercase_digest = (
        manifest_digest.split(":", 1)[0]
        + ":"
        + manifest_digest.split(":", 1)[1].upper()
    )

    def capture_request(url, *_args, **_kwargs):
        requested.append(url)
        return Response(manifest_body)

    puller.make_request = capture_request
    assert puller.get_manifest("team/app", uppercase_digest, None) == manifest
    assert requested == [
        f"https://registry-1.docker.io/v2/team/app/manifests/{manifest_digest}"
    ]

    def selected_manifest(body, digest, size):
        index = json.dumps(
            {
                "schemaVersion": 2,
                "manifests": [
                    {
                        "digest": digest,
                        "size": size,
                        "platform": {"os": "linux", "architecture": "amd64"},
                    }
                ],
            }
        ).encode()
        selected_puller = DockerImagePuller()
        responses = iter([Response(index), Response(body)])
        selected_puller.make_request = lambda *_args, **_kwargs: next(responses)
        return selected_puller

    with pytest.raises(RegistryError, match="digest mismatch"):
        selected_manifest(
            manifest_body, "sha256:" + "0" * 64, len(manifest_body)
        ).get_manifest("team/app", "latest", None)
    with pytest.raises(RegistryError, match="size mismatch"):
        selected_manifest(
            manifest_body, manifest_digest, len(manifest_body) + 1
        ).get_manifest("team/app", "latest", None)
    with pytest.raises(RegistryError, match="exceeds the 1 B limit"):
        selected_manifest(manifest_body, manifest_digest, 1).get_manifest(
            "team/app", "latest", None
        )

    nested_body = json.dumps({"schemaVersion": 2, "manifests": []}).encode()
    with pytest.raises(RegistryError, match="Nested image indexes"):
        selected_manifest(
            nested_body, content_digest(nested_body), len(nested_body)
        ).get_manifest("team/app", "latest", None)

    schema_v1 = json.dumps({"schemaVersion": 1, "fsLayers": []}).encode()
    puller.make_request = lambda *_args, **_kwargs: Response(schema_v1)
    with pytest.raises(RegistryError, match="Schema v1"):
        puller.get_manifest("team/app", "latest", None)

    schema_v1_string = json.dumps({"schemaVersion": "1", "fsLayers": []}).encode()
    puller.make_request = lambda *_args, **_kwargs: Response(schema_v1_string)
    with pytest.raises(RegistryError, match="Schema v1"):
        puller.get_manifest("team/app", "latest", None)


def test_manifest_response_size_is_bounded():
    puller = DockerImagePuller()
    puller.make_request = lambda *_args, **_kwargs: Response(
        b"x" * (docker_pull.MAX_MANIFEST_SIZE + 1)
    )
    with pytest.raises(RegistryError, match="exceeds the 4.0 MB limit"):
        puller.get_manifest("team/app", "latest", None)


def test_manifest_401_refreshes_authentication():
    manifest = {"schemaVersion": 2, "config": {}, "layers": []}
    body = json.dumps(manifest).encode()
    response_headers = Message()
    response_headers["WWW-Authenticate"] = 'Bearer realm="https://login.example/token"'
    requests = []
    puller = DockerImagePuller()

    def make_request(url, headers=None):
        requests.append(headers)
        if len(requests) == 1:
            raise HTTPError(url, 401, "Unauthorized", response_headers, None)
        return Response(body)

    puller.make_request = make_request
    puller.get_auth_header = lambda *_args, **_kwargs: "Bearer fresh"
    assert puller.get_manifest("team/app", "latest", "Bearer stale") == manifest
    assert [headers["Authorization"] for headers in requests] == [
        "Bearer stale",
        "Bearer fresh",
    ]


@pytest.mark.parametrize(
    ("config_architecture", "variant", "requested_architecture"),
    [
        ("aarch64", None, "arm64"),
        ("x86_64", None, "amd64"),
        ("i386", None, "386"),
        ("arm", "v8", "arm64"),
    ],
)
def test_config_architecture_aliases(
    config_architecture, variant, requested_architecture
):
    config_data = json.dumps(
        {"architecture": config_architecture, "variant": variant, "os": "linux"}
    ).encode()
    config_digest = content_digest(config_data)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [{"digest": "sha256:" + "1" * 64, "size": 1}],
    }
    with pytest.raises(RegistryError, match="Failed to download layer"):
        offline_puller(manifest, config_data).pull_image(
            "registry.example/team/app:latest", architecture=requested_architecture
        )


def test_config_variant_and_os_comparison_is_case_insensitive():
    config_data = json.dumps(
        {"architecture": "arm", "variant": "V7", "os": "Linux"}
    ).encode()
    config_digest = content_digest(config_data)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [{"digest": "sha256:" + "1" * 64, "size": 1}],
    }
    with pytest.raises(RegistryError, match="Failed to download layer"):
        offline_puller(manifest, config_data).pull_image(
            "registry.example/team/app:latest", architecture="arm", variant="v7"
        )


def test_partial_pull_is_atomic(tmp_path):
    config_data = b'{"architecture":"amd64","os":"linux"}'
    config_digest = content_digest(config_data)
    first_layer = b"first layer"
    first_digest = content_digest(first_layer)
    second_layer = b"second layer"
    second_digest = content_digest(second_layer)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [
            {"digest": first_digest, "size": len(first_layer)},
            {"digest": second_digest, "size": len(second_layer)},
        ],
    }
    puller = DockerImagePuller()
    puller.get_auth_header = lambda image_name, challenge=None: None
    puller.get_manifest = lambda *_args, **_kwargs: manifest

    def download_blob(
        image_name,
        digest,
        authorization,
        retry_with_new_auth=True,
        expected_size=None,
        destination=None,
    ):
        if digest == config_digest:
            return config_data
        if digest == first_digest and destination:
            with open(destination, "wb") as blob_file:
                blob_file.write(first_layer)
            return destination
        return None

    puller.download_blob = download_blob
    output = tmp_path / "image.tar"
    output.write_bytes(b"existing")
    with pytest.raises(RegistryError, match="Failed to download layer"):
        puller.pull_image("registry.example/team/app:latest", str(output))
    assert output.read_bytes() == b"existing"


def test_missing_output_directory_fails_before_network(tmp_path):
    puller = DockerImagePuller()
    with patch.object(puller, "get_auth_header") as get_auth_header:
        with pytest.raises(ValueError, match="Output directory does not exist"):
            puller.pull_image(
                "registry.example/team/app:latest",
                str(tmp_path / "missing" / "image.tar"),
            )
    get_auth_header.assert_not_called()


@pytest.mark.parametrize(
    ("media_type", "error"),
    [
        (
            "application/vnd.docker.image.rootfs.foreign.diff.tar.gzip",
            "foreign/nondistributable",
        ),
        ("application/vnd.oci.image.layer.v1.tar+zstd", "Zstandard"),
    ],
)
def test_unsupported_layers_fail_before_layer_download(media_type, error):
    config_data = b'{"architecture":"amd64","os":"linux"}'
    config_digest = content_digest(config_data)
    layer_digest = "sha256:" + "1" * 64
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [
            {"digest": layer_digest, "size": None, "mediaType": None},
            {"digest": layer_digest, "size": 1, "mediaType": media_type},
        ],
    }
    downloaded = []
    puller = offline_puller(manifest, config_data, downloaded)
    with pytest.raises(RegistryError, match=error):
        puller.pull_image("registry.example/team/app:latest")
    assert downloaded == [config_digest]


def test_malformed_layer_descriptor_is_a_registry_error():
    config_data = b'{"architecture":"amd64","os":"linux"}'
    config_digest = content_digest(config_data)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [{"digest": "sha256:" + "1" * 64, "size": "unknown"}],
    }
    with pytest.raises(RegistryError, match="invalid size"):
        offline_puller(manifest, config_data).pull_image(
            "registry.example/team/app:latest"
        )

    puller = DockerImagePuller()
    puller.get_auth_header = lambda *_args, **_kwargs: None
    puller.get_manifest = lambda *_args, **_kwargs: {"config": "invalid"}
    with pytest.raises(RegistryError, match="config descriptor is not an object"):
        puller.pull_image("registry.example/team/app:latest")


def test_duplicate_layers_are_downloaded_once(tmp_path):
    layer_data = b"layer data"
    layer_digest = content_digest(layer_data)
    config_data = json.dumps(
        {
            "architecture": "amd64",
            "os": "linux",
            "rootfs": {"diff_ids": [layer_digest, layer_digest]},
        }
    ).encode()
    config_digest = content_digest(config_data)
    layer = {"digest": layer_digest, "size": len(layer_data)}
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [layer, layer],
    }
    downloaded = []
    puller = DockerImagePuller()
    puller.get_auth_header = lambda *_args, **_kwargs: None
    puller.get_manifest = lambda *_args, **_kwargs: manifest

    def download_blob(
        image_name,
        digest,
        authorization,
        retry_with_new_auth=True,
        expected_size=None,
        destination=None,
    ):
        del image_name, authorization, retry_with_new_auth, expected_size
        downloaded.append(digest)
        if digest == config_digest:
            return config_data
        assert destination is not None
        with open(destination, "wb") as blob_file:
            blob_file.write(layer_data)
        return destination

    puller.download_blob = download_blob
    puller.pull_image("registry.example/team/app:latest", str(tmp_path / "image.tar"))
    assert downloaded == [config_digest, layer_digest]


def test_duplicate_layers_with_conflicting_sizes_fail(tmp_path):
    layer_data = b"layer data"
    layer_digest = content_digest(layer_data)
    config_data = b'{"architecture":"amd64","os":"linux"}'
    config_digest = content_digest(config_data)
    manifest = {
        "config": {"digest": config_digest, "size": len(config_data)},
        "layers": [
            {"digest": layer_digest, "size": len(layer_data)},
            {"digest": layer_digest, "size": len(layer_data) + 1},
        ],
    }
    downloaded = []
    puller = DockerImagePuller()
    puller.get_auth_header = lambda *_args, **_kwargs: None
    puller.get_manifest = lambda *_args, **_kwargs: manifest

    def download_blob(
        image_name,
        digest,
        authorization,
        retry_with_new_auth=True,
        expected_size=None,
        destination=None,
    ):
        del image_name, authorization, retry_with_new_auth, expected_size
        downloaded.append(digest)
        if digest == config_digest:
            return config_data
        assert destination is not None
        with open(destination, "wb") as blob_file:
            blob_file.write(layer_data)
        return destination

    puller.download_blob = download_blob
    with pytest.raises(RegistryError, match="inconsistent sizes"):
        puller.pull_image(
            "registry.example/team/app:latest", str(tmp_path / "image.tar")
        )
    assert downloaded == [config_digest, layer_digest]


def test_modern_archive_metadata(tmp_path):
    layer_data = b"layer tar data"
    layer_digest = content_digest(layer_data)
    config_data = json.dumps(
        {"rootfs": {"diff_ids": [layer_digest, layer_digest]}}
    ).encode()
    config_digest = content_digest(config_data)
    layer_path = tmp_path / "layer.blob"
    layer_path.write_bytes(layer_data)
    manifest = {"config": {"digest": config_digest}}
    layer = {
        "digest": layer_digest,
        "size": len(layer_data),
        "path": str(layer_path),
    }
    puller = DockerImagePuller()

    tagged_output = tmp_path / "tagged.tar"
    puller.create_docker_tar(
        "registry.example/team/app",
        "v1",
        manifest,
        config_data,
        json.loads(config_data),
        [layer, layer],
        str(tagged_output),
    )
    with tarfile.open(tagged_output) as archive:
        manifest_file = archive.extractfile("manifest.json")
        assert manifest_file is not None
        archive_manifest = json.load(manifest_file)
        archive_names = archive.getnames()
    layer_name = f"{layer_digest.split(':', 1)[1]}/layer.tar"
    assert archive_manifest[0]["RepoTags"] == ["registry.example/team/app:v1"]
    assert archive_manifest[0]["Layers"] == [layer_name, layer_name]
    assert archive_names.count(layer_name) == 1
    assert "repositories" not in archive_names
    assert not any(
        name.endswith("/VERSION") or name.endswith("/json") for name in archive_names
    )

    config_data = json.dumps({"rootfs": {"diff_ids": [layer_digest]}}).encode()
    manifest = {"config": {"digest": content_digest(config_data)}}
    digest_output = tmp_path / "digest.tar"
    puller.create_docker_tar(
        "registry.example/team/app",
        None,
        manifest,
        config_data,
        json.loads(config_data),
        [layer],
        str(digest_output),
    )
    with tarfile.open(digest_output) as archive:
        manifest_file = archive.extractfile("manifest.json")
        assert manifest_file is not None
        archive_manifest = json.load(manifest_file)
    assert archive_manifest[0]["RepoTags"] is None


def test_gzip_layer_is_unpacked_in_archive(tmp_path):
    layer_data = b"uncompressed layer tar payload"
    compressed = gzip.compress(layer_data)
    layer_digest = content_digest(compressed)
    diff_id = content_digest(layer_data)
    config_data = json.dumps({"rootfs": {"diff_ids": [diff_id]}}).encode()
    config_digest = content_digest(config_data)
    layer_path = tmp_path / "layer.blob"
    layer_path.write_bytes(compressed)
    output = tmp_path / "image.tar"

    DockerImagePuller().create_docker_tar(
        "registry.example/team/app",
        "v1",
        {"config": {"digest": config_digest}},
        config_data,
        json.loads(config_data),
        [
            {
                "digest": layer_digest,
                "size": len(compressed),
                "path": str(layer_path),
            }
        ],
        str(output),
    )

    layer_name = f"{layer_digest.split(':', 1)[1]}/layer.tar"
    with tarfile.open(output) as archive:
        layer_file = archive.extractfile(layer_name)
        assert layer_file is not None
        archived_layer = layer_file.read()
    assert archived_layer == layer_data
    assert content_digest(archived_layer) == diff_id


@pytest.mark.parametrize(
    ("diff_ids", "error"),
    [([], "count does not match"), (["sha256:" + "0" * 64], "diff-ID mismatch")],
)
def test_layer_diff_ids_are_verified(tmp_path, diff_ids, error):
    layer_data = b"layer tar data"
    layer_path = tmp_path / "layer.blob"
    layer_path.write_bytes(layer_data)
    config_data = json.dumps({"rootfs": {"diff_ids": diff_ids}}).encode()
    layer = {
        "digest": content_digest(layer_data),
        "size": len(layer_data),
        "path": str(layer_path),
    }
    with pytest.raises(RegistryError, match=error):
        DockerImagePuller().create_docker_tar(
            "registry.example/team/app",
            "v1",
            {"config": {"digest": content_digest(config_data)}},
            config_data,
            json.loads(config_data),
            [layer],
            str(tmp_path / "image.tar"),
        )


def test_zstandard_magic_is_rejected_during_archive_creation(tmp_path):
    config_data = json.dumps(
        {"rootfs": {"diff_ids": [content_digest(b"unused")]}}
    ).encode()
    config_digest = content_digest(config_data)
    layer_path = tmp_path / "layer.blob"
    layer_path.write_bytes(b"\x28\xb5\x2f\xfdpayload")
    layer = {
        "digest": content_digest(layer_path.read_bytes()),
        "size": layer_path.stat().st_size,
        "path": str(layer_path),
    }
    with pytest.raises(RegistryError, match="Zstandard"):
        DockerImagePuller().create_docker_tar(
            "registry.example/team/app",
            "v1",
            {"config": {"digest": config_digest}},
            config_data,
            json.loads(config_data),
            [layer],
            str(tmp_path / "image.tar"),
        )


@pytest.mark.parametrize(
    ("arguments", "input_text"),
    [
        (["image", "--token", "x", "--username", "u", "--password", "p"], None),
        (["image", "--plain-http", "--insecure"], None),
        (["image", "--token-stdin"], ""),
        (["image", "--token-stdin"], "line1\nline2\n"),
    ],
)
def test_cli_validation(arguments, input_text):
    script = os.path.join(os.path.dirname(__file__), "docker_pull.py")
    result = subprocess.run(
        [sys.executable, script, *arguments],
        input=input_text,
        text=True,
        capture_output=True,
        timeout=10,
        check=False,
    )
    assert result.returncode == 2


def test_malformed_proxy_is_reported_without_a_traceback():
    script = os.path.join(os.path.dirname(__file__), "docker_pull.py")
    result = subprocess.run(
        [
            sys.executable,
            script,
            "image",
            "--proxy",
            "http://proxy.example:invalid",
            "--proxy-auth",
            "user:secret",
        ],
        text=True,
        capture_output=True,
        timeout=10,
        check=False,
    )
    assert result.returncode == 1
    assert "Invalid http_proxy URL" in result.stdout
    assert "Traceback" not in result.stderr


def test_malformed_registry_is_reported_without_a_traceback():
    script = os.path.join(os.path.dirname(__file__), "docker_pull.py")
    result = subprocess.run(
        [sys.executable, script, "registry.example:invalid/team/app:latest"],
        text=True,
        capture_output=True,
        timeout=10,
        check=False,
    )
    assert result.returncode == 1
    assert "Invalid registry host" in result.stdout
    assert "Traceback" not in result.stderr


def test_cli_passes_token_and_variant_to_the_puller(monkeypatch):
    monkeypatch.setattr(
        sys,
        "argv",
        [
            "docker_pull.py",
            "registry.example/team/app:latest",
            "--token",
            "secret",
            "--arch",
            "arm",
            "--variant",
            "v7",
        ],
    )
    with patch.object(docker_pull, "DockerImagePuller") as puller_type:
        docker_pull.main()
    puller_type.assert_called_once_with(
        auth_token="secret", proxy_config={}, username=None, password=None
    )
    puller_type.return_value.pull_image.assert_called_once_with(
        "registry.example/team/app:latest",
        None,
        architecture="arm",
        os_type="linux",
        variant="v7",
        plain_http=False,
    )


@pytest.mark.parametrize(
    ("arguments", "shows_info"),
    [
        (["-v", "--quiet"], False),
        (["-v", "--log-level", "ERROR"], False),
        (["--quiet", "--log-level", "DEBUG"], False),
        (["--debug", "--quiet"], True),
    ],
)
def test_cli_logging_precedence(arguments, shows_info):
    script = os.path.join(os.path.dirname(__file__), "docker_pull.py")
    result = subprocess.run(
        [
            sys.executable,
            script,
            "127.0.0.1:1/team/app:latest",
            "--plain-http",
            *arguments,
        ],
        text=True,
        capture_output=True,
        timeout=10,
        check=False,
    )
    assert result.returncode == 1
    assert ("Downloading image:" in result.stdout) is shows_info