From 5e8ee1ca7ddfba1af89be6147fbfca1fa7b75f07 Mon Sep 17 00:00:00 2001
From: Simon Li <orpheus+devel@gmail.com>
Date: Sun, 12 Feb 2023 00:16:14 +0000
Subject: [PATCH] Add registry optional credentials to push()

---
 repo2docker/docker.py     |  2 ++
 repo2docker/engine.py     | 37 +++++++++++++++++++++++++++++++++
 tests/unit/test_docker.py | 43 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+)

diff --git a/repo2docker/docker.py b/repo2docker/docker.py
index 39d9e044..170eea3f 100644
--- a/repo2docker/docker.py
+++ b/repo2docker/docker.py
@@ -120,6 +120,8 @@ class DockerEngine(ContainerEngine):
         return Image(tags=image["RepoTags"], config=image["ContainerConfig"])
 
     def push(self, image_spec):
+        if self.registry_credentials:
+            self._apiclient.login(**self.registry_credentials)
         return self._apiclient.push(image_spec, stream=True)
 
     def run(
diff --git a/repo2docker/engine.py b/repo2docker/engine.py
index f3febccc..a8897a64 100644
--- a/repo2docker/engine.py
+++ b/repo2docker/engine.py
@@ -2,8 +2,11 @@
 Interface for a repo2docker container engine
 """
 
+import json
+import os
 from abc import ABC, abstractmethod
 
+from traitlets import Dict, default
 from traitlets.config import LoggingConfigurable
 
 # Based on https://docker-py.readthedocs.io/en/4.2.0/containers.html
@@ -142,6 +145,37 @@ class ContainerEngine(LoggingConfigurable):
     Initialised with a reference to the parent so can also be configured using traitlets.
     """
 
+    registry_credentials = Dict(
+        help="""
+        Credentials dictionary, if set will be used to authenticate with
+        the registry. Typically this will include the keys:
+
+            - `username`: The registry username
+            - `password`: The registry password or token
+            - `registry`: The registry URL
+
+        This can also be set by passing a JSON object in the
+        CONTAINER_ENGINE_REGISTRY_CREDENTIALS environment variable.
+        """,
+        config=True,
+    )
+
+    @default("registry_credentials")
+    def _registry_credentials_default(self):
+        """
+        Set the registry credentials from CONTAINER_ENGINE_REGISTRY_CREDENTIALS
+        """
+        obj = os.getenv("CONTAINER_ENGINE_REGISTRY_CREDENTIALS")
+        if obj:
+            try:
+                return json.loads(obj)
+            except json.JSONDecodeError:
+                self.log.error(
+                    "CONTAINER_ENGINE_REGISTRY_CREDENTIALS is not valid JSON"
+                )
+                raise
+        return {}
+
     string_output = True
     """
     Whether progress events should be strings or an object.
@@ -251,6 +285,9 @@ class ContainerEngine(LoggingConfigurable):
         """
         Push image to a registry
 
+        If the registry_credentials traitlets is set it should be used to
+        authenticate with the registry before pushing.
+
         Parameters
         ----------
         image_spec : str
diff --git a/tests/unit/test_docker.py b/tests/unit/test_docker.py
index 2fe0b6f5..c47aa4fa 100644
--- a/tests/unit/test_docker.py
+++ b/tests/unit/test_docker.py
@@ -2,6 +2,9 @@
 
 import os
 from subprocess import check_output
+from unittest.mock import Mock, patch
+
+from repo2docker.docker import DockerEngine
 
 repo_root = os.path.abspath(
     os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
@@ -19,3 +22,43 @@ def test_git_credential_env():
         .strip()
     )
     assert out == credential_env
+
+
+class MockDockerEngine(DockerEngine):
+    def __init__(self, *args, **kwargs):
+        self._apiclient = Mock()
+
+
+def test_docker_push_no_credentials():
+    engine = MockDockerEngine()
+
+    engine.push("image")
+
+    assert len(engine._apiclient.method_calls) == 1
+    engine._apiclient.push.assert_called_once_with("image", stream=True)
+
+
+def test_docker_push_dict_credentials():
+    engine = MockDockerEngine()
+    engine.registry_credentials = {"username": "abc", "password": "def"}
+
+    engine.push("image")
+
+    assert len(engine._apiclient.method_calls) == 2
+    engine._apiclient.login.assert_called_once_with(username="abc", password="def")
+    engine._apiclient.push.assert_called_once_with("image", stream=True)
+
+
+def test_docker_push_env_credentials():
+    engine = MockDockerEngine()
+    with patch.dict(
+        "os.environ",
+        {
+            "CONTAINER_ENGINE_REGISTRY_CREDENTIALS": '{"username": "abc", "password": "def"}'
+        },
+    ):
+        engine.push("image")
+
+    assert len(engine._apiclient.method_calls) == 2
+    engine._apiclient.login.assert_called_once_with(username="abc", password="def")
+    engine._apiclient.push.assert_called_once_with("image", stream=True)