feat(core): Private registry (#566)

Ref #562 This enhancement adds capability to utilize the env var `DOCKER_AUTH_CONFIG` in-order to login to a private docker registry. --------- Co-authored-by: David Ankin <[email protected]>
testcontainers · May 25, 2024 · 59fbcfa · 59fbcfa
1 parent 9d2ceb6
commit 59fbcfa
Show file tree

Hide file tree

Showing 5 changed files with 139 additions and 2 deletions.
diff --git a/core/testcontainers/core/config.py b/core/testcontainers/core/config.py
@@ -1,8 +1,9 @@
 from dataclasses import dataclass, field
+from logging import warning
 from os import environ
 from os.path import exists
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 MAX_TRIES = int(environ.get("TC_MAX_TRIES", 120))
 SLEEP_TIME = int(environ.get("TC_POOLING_INTERVAL", 1))
@@ -37,6 +38,9 @@ def read_tc_properties() -> dict[str, str]:
     return settings
 
 
+_WARNINGS = {"DOCKER_AUTH_CONFIG": "DOCKER_AUTH_CONFIG is experimental, see testcontainers/testcontainers-python#566"}
+
+
 @dataclass
 class TestcontainersConfiguration:
     max_tries: int = MAX_TRIES
@@ -47,6 +51,19 @@ class TestcontainersConfiguration:
     ryuk_docker_socket: str = RYUK_DOCKER_SOCKET
     ryuk_reconnection_timeout: str = RYUK_RECONNECTION_TIMEOUT
     tc_properties: dict[str, str] = field(default_factory=read_tc_properties)
+    _docker_auth_config: Optional[str] = field(default_factory=lambda: environ.get("DOCKER_AUTH_CONFIG"))
+
+    @property
+    def docker_auth_config(self):
+        if "DOCKER_AUTH_CONFIG" in _WARNINGS:
+            warning(_WARNINGS.pop("DOCKER_AUTH_CONFIG"))
+        return self._docker_auth_config
+
+    @docker_auth_config.setter
+    def docker_auth_config(self, value: str):
+        if "DOCKER_AUTH_CONFIG" in _WARNINGS:
+            warning(_WARNINGS.pop("DOCKER_AUTH_CONFIG"))
+        self._docker_auth_config = value
 
     def tc_properties_get_tc_host(self) -> Union[str, None]:
         return self.tc_properties.get("tc.host")

diff --git a/core/testcontainers/core/docker_client.py b/core/testcontainers/core/docker_client.py
@@ -24,7 +24,7 @@
 
 from testcontainers.core.config import testcontainers_config as c
 from testcontainers.core.labels import SESSION_ID, create_labels
-from testcontainers.core.utils import default_gateway_ip, inside_container, setup_logger
+from testcontainers.core.utils import default_gateway_ip, inside_container, parse_docker_auth_config, setup_logger
 
 LOGGER = setup_logger(__name__)
 
@@ -57,6 +57,9 @@ def __init__(self, **kwargs) -> None:
         self.client.api.headers["x-tc-sid"] = SESSION_ID
         self.client.api.headers["User-Agent"] = "tc-python/" + importlib.metadata.version("testcontainers")
 
+        if docker_auth_config := get_docker_auth_config():
+            self.login(docker_auth_config)
+
     @_wrapped_container_collection
     def run(
         self,
@@ -183,6 +186,18 @@ def host(self) -> str:
                 return ip_address
         return "localhost"
 
+    def login(self, docker_auth_config: str) -> None:
+        """
+        Login to a docker registry using the given auth config.
+        """
+        auth_config = parse_docker_auth_config(docker_auth_config)[0]  # Only using the first auth config
+        login_info = self.client.login(**auth_config._asdict())
+        LOGGER.debug(f"logged in using {login_info}")
+
 
 def get_docker_host() -> Optional[str]:
     return c.tc_properties_get_tc_host() or os.getenv("DOCKER_HOST")
+
+
+def get_docker_auth_config() -> Optional[str]:
+    return c.docker_auth_config
diff --git a/core/testcontainers/core/utils.py b/core/testcontainers/core/utils.py
@@ -1,13 +1,18 @@
+import base64
+import json
 import logging
 import os
 import platform
 import subprocess
 import sys
+from collections import namedtuple
 
 LINUX = "linux"
 MAC = "mac"
 WIN = "win"
 
+DockerAuthInfo = namedtuple("DockerAuthInfo", ["registry", "username", "password"])
+
 
 def setup_logger(name: str) -> logging.Logger:
     logger = logging.getLogger(name)
@@ -77,3 +82,29 @@ def raise_for_deprecated_parameter(kwargs: dict, name: str, replacement: str) ->
     if kwargs.pop(name, None):
         raise ValueError(f"Use `{replacement}` instead of `{name}`")
     return kwargs
+
+
+def parse_docker_auth_config(auth_config: str) -> list[DockerAuthInfo]:
+    """
+    Parse the docker auth config from a string.
+
+    Example:
+    {
+        "auths": {
+            "https://index.docker.io/v1/": {
+                "auth": "dXNlcm5hbWU6cGFzc3dvcmQ="
+            }
+        }
+    }
+    """
+    auth_info: list[DockerAuthInfo] = []
+    try:
+        auth_config_dict: dict = json.loads(auth_config).get("auths")
+        for registry, auth in auth_config_dict.items():
+            auth_str = auth.get("auth")
+            auth_str = base64.b64decode(auth_str).decode("utf-8")
+            username, password = auth_str.split(":")
+            auth_info.append(DockerAuthInfo(registry, username, password))
+        return auth_info
+    except (json.JSONDecodeError, KeyError, ValueError) as exp:
+        raise ValueError("Could not parse docker auth config") from exp
diff --git a/core/tests/test_docker_client.py b/core/tests/test_docker_client.py
@@ -1,9 +1,14 @@
+import os
+from collections import namedtuple
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
 import docker
 
+from testcontainers.core.config import testcontainers_config as c
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.docker_client import DockerClient
+from testcontainers.core.utils import parse_docker_auth_config
 
 
 def test_docker_client_from_env():
@@ -15,6 +20,33 @@ def test_docker_client_from_env():
     mock_docker.from_env.assert_called_with(**test_kwargs)
 
 
+def test_docker_client_login_no_login():
+    with patch.dict(os.environ, {}, clear=True):
+        mock_docker = MagicMock(spec=docker)
+        with patch("testcontainers.core.docker_client.docker", mock_docker):
+            DockerClient()
+
+        mock_docker.from_env.return_value.login.assert_not_called()
+
+
+def test_docker_client_login():
+    mock_docker = MagicMock(spec=docker)
+    mock_parse_docker_auth_config = MagicMock(spec=parse_docker_auth_config)
+    mock_utils = MagicMock()
+    mock_utils.parse_docker_auth_config = mock_parse_docker_auth_config
+    TestAuth = namedtuple("Auth", "value")
+    mock_parse_docker_auth_config.return_value = [TestAuth("test")]
+
+    with (
+        mock.patch.object(c, "_docker_auth_config", "test"),
+        patch("testcontainers.core.docker_client.docker", mock_docker),
+        patch("testcontainers.core.docker_client.parse_docker_auth_config", mock_parse_docker_auth_config),
+    ):
+        DockerClient()
+
+    mock_docker.from_env.return_value.login.assert_called_with(**{"value": "test"})
+
+
 def test_container_docker_client_kw():
     test_kwargs = {"test_kw": "test_value"}
     mock_docker = MagicMock(spec=docker)

diff --git a/core/tests/test_utils.py b/core/tests/test_utils.py
@@ -0,0 +1,42 @@
+import json
+
+from testcontainers.core.utils import parse_docker_auth_config, DockerAuthInfo
+
+
+def test_parse_docker_auth_config():
+    auth_config_json = '{"auths":{"https://index.docker.io/v1/":{"auth":"dXNlcm5hbWU6cGFzc3dvcmQ="}}}'
+    auth_info = parse_docker_auth_config(auth_config_json)
+    assert len(auth_info) == 1
+    assert auth_info[0] == DockerAuthInfo(
+        registry="https://index.docker.io/v1/",
+        username="username",
+        password="password",
+    )
+
+
+def test_parse_docker_auth_config_multiple():
+    auth_dict = {
+        "auths": {
+            "localhost:5000": {"auth": "dXNlcjE6cGFzczE=="},
+            "https://example.com": {"auth": "dXNlcl9uZXc6cGFzc19uZXc=="},
+            "example2.com": {"auth": "YWJjOjEyMw==="},
+        }
+    }
+    auth_config_json = json.dumps(auth_dict)
+    auth_info = parse_docker_auth_config(auth_config_json)
+    assert len(auth_info) == 3
+    assert auth_info[0] == DockerAuthInfo(
+        registry="localhost:5000",
+        username="user1",
+        password="pass1",
+    )
+    assert auth_info[1] == DockerAuthInfo(
+        registry="https://example.com",
+        username="user_new",
+        password="pass_new",
+    )
+    assert auth_info[2] == DockerAuthInfo(
+        registry="example2.com",
+        username="abc",
+        password="123",
+    )