From f7308104012c8ed9c97840860a2e27f258e7377b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AE=B5=E9=80=B8?= <13691419520@163.com>
Date: Sun, 15 Nov 2020 12:34:42 -0800
Subject: [PATCH] Allow override DOCKER_HUB_REPO

---
 bugswarm/client/bugswarm.py |  7 +++++--
 bugswarm/client/docker.py   | 30 ++++++++++++++++++------------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/bugswarm/client/bugswarm.py b/bugswarm/client/bugswarm.py
index ecc19dc..46a2bb0 100644
--- a/bugswarm/client/bugswarm.py
+++ b/bugswarm/client/bugswarm.py
@@ -20,6 +20,9 @@ def cli():
 
 
 @cli.command(cls=MyCommand)
+@click.option('--docker-hub-repo', default=None,
+              type=str,
+              help='Override the default Docker Hub repo (bugswarm/images)')
 @click.option('--image-tag', required=True,
               type=str,
               help='The artifact image tag.')
@@ -33,12 +36,12 @@ def cli():
               help='If enabled, artifact containers will be cleaned up automatically after use. '
                    'Disable this behavior if you want to inspect the container filesystem after use. '
                    'Enabled by default.')
-def run(image_tag, use_sandbox, pipe_stdin, rm):
+def run(docker_hub_repo, image_tag, use_sandbox, pipe_stdin, rm):
     """Start an artifact container."""
     # If the script does not already have sudo privileges, then explain to the user why the password prompt will appear.
     if os.getuid() != 0:
         log.info('Docker requires sudo privileges.')
-    docker.docker_run(image_tag, use_sandbox, pipe_stdin, rm)
+    docker.docker_run(docker_hub_repo, image_tag, use_sandbox, pipe_stdin, rm)
 
 
 @cli.command(cls=MyCommand)
diff --git a/bugswarm/client/docker.py b/bugswarm/client/docker.py
index 9225785..769b545 100644
--- a/bugswarm/client/docker.py
+++ b/bugswarm/client/docker.py
@@ -13,14 +13,15 @@
 
 # By default, this function downloads the image, enters the container, and executes '/bin/bash' in the container.
 # The executed script can be changed by passing the script argument.
-def docker_run(image_tag, use_sandbox, use_pipe_stdin, use_rm):
+def docker_run(docker_hub_repo, image_tag, use_sandbox, use_pipe_stdin, use_rm):
+    assert (isinstance(docker_hub_repo, str) and not docker_hub_repo.isspace()) or docker_hub_repo is None
     assert isinstance(image_tag, str) and not image_tag.isspace()
     assert isinstance(use_sandbox, bool)
     assert isinstance(use_pipe_stdin, bool)
     assert isinstance(use_rm, bool)
 
     # First, try to pull the image.
-    ok = docker_pull(image_tag)
+    ok = docker_pull(docker_hub_repo, image_tag)
     if not ok:
         return False
 
@@ -42,7 +43,7 @@ def docker_run(image_tag, use_sandbox, use_pipe_stdin, use_rm):
     if use_rm:
         log.info('The container will be cleaned up after use.')
 
-    image_location = _image_location(image_tag)
+    image_location = _image_location(docker_hub_repo, image_tag)
 
     # Prepare the arguments for the docker run command.
     volume_args = ['-v', '{}:{}'.format(host_sandbox, container_sandbox)] if use_sandbox else []
@@ -75,15 +76,16 @@ def docker_run(image_tag, use_sandbox, use_pipe_stdin, use_rm):
     return returncode == 0
 
 
-def docker_pull(image_tag):
+def docker_pull(docker_hub_repo, image_tag):
+    assert (isinstance(docker_hub_repo, str) and not docker_hub_repo.isspace()) or docker_hub_repo is None
     assert image_tag
     assert isinstance(image_tag, str)
 
     # Exit early if the image already exists locally.
-    if _image_exists_locally(image_tag):
+    if _image_exists_locally(docker_hub_repo, image_tag):
         return True
 
-    image_location = _image_location(image_tag)
+    image_location = _image_location(docker_hub_repo, image_tag)
     command = 'sudo docker pull {}'.format(image_location)
     _, _, returncode = ShellWrapper.run_commands(command, shell=True)
     if returncode != 0:
@@ -94,8 +96,8 @@ def docker_pull(image_tag):
 
 
 # Returns True if the image already exists locally.
-def _docker_image_inspect(image_tag):
-    image_location = _image_location(image_tag)
+def _docker_image_inspect(docker_hub_repo, image_tag):
+    image_location = _image_location(docker_hub_repo, image_tag)
     command = 'sudo docker image inspect {}'.format(image_location)
     _, _, returncode = ShellWrapper.run_commands(command,
                                                  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True)
@@ -106,14 +108,18 @@ def _docker_image_inspect(image_tag):
 
 
 # Returns True if the image already exists locally.
-def _image_exists_locally(image_tag):
-    return _docker_image_inspect(image_tag)
+def _image_exists_locally(docker_hub_repo, image_tag):
+    return _docker_image_inspect(docker_hub_repo, image_tag)
 
 
-def _image_location(image_tag):
+def _image_location(docker_hub_repo, image_tag):
+    assert (isinstance(docker_hub_repo, str) and not docker_hub_repo.isspace()) or docker_hub_repo is None
     assert image_tag
     assert isinstance(image_tag, str)
-    return DOCKER_HUB_REPO + ':' + image_tag
+    if docker_hub_repo is None:
+        return DOCKER_HUB_REPO + ':' + image_tag
+    else:
+        return docker_hub_repo + ':' + image_tag
 
 
 def _default_host_sandbox():