Allow users to override TUF repository URL prefix (#3660)

* allow users to override TUF repository URL prefix

* Apply suggestions from code review

Co-Authored-By: zippolyte <zippolyte@users.noreply.github.com>

* black

datadog_checks_downloader/datadog_checks/downloader/cli.py

@@ -11,7 +11,8 @@
 from tuf.exceptions import UnknownTargetError
 
 # 2nd party.
-from .download import TUFDownloader
+# 2nd party.
+from .download import REPOSITORY_URL_PREFIX, TUFDownloader
 from .exceptions import NonCanonicalVersion, NonDatadogPackage, NoSuchDatadogPackageOrVersion
 
 # Private module functions.
@@ -41,13 +42,18 @@ def download():
         'standard_distribution_name', type=str, help='Standard distribution name of the desired Datadog check.'
     )
 
+    parser.add_argument(
+        '--repository', type=str, default=REPOSITORY_URL_PREFIX, help='The complete URL prefix for the TUF repository.'
+    )
+
     parser.add_argument('--version', type=str, default=None, help='The version number of the desired Datadog check.')
 
     parser.add_argument(
         '-v', '--verbose', action='count', default=0, help='Show verbose information about TUF and in-toto.'
     )
 
     args = parser.parse_args()
+    repository_url_prefix = args.repository
     standard_distribution_name = args.standard_distribution_name
     version = args.version
     verbose = args.verbose
@@ -56,7 +62,7 @@ def download():
         raise NonDatadogPackage(standard_distribution_name)
     else:
         wheel_distribution_name = __get_wheel_distribution_name(standard_distribution_name)
-        tuf_downloader = TUFDownloader(verbose=verbose)
+        tuf_downloader = TUFDownloader(repository_url_prefix=repository_url_prefix, verbose=verbose)
 
         if not version:
             version = tuf_downloader.get_latest_version(standard_distribution_name, wheel_distribution_name)

datadog_checks_downloader/datadog_checks/downloader/download.py

@@ -46,14 +46,15 @@
         'confined_target_dirs': [''],
     }
 }
+REPOSITORY_URL_PREFIX = 'https://dd-integrations-core-wheels-build-stable.datadoghq.com'
 
 
 # Global variables.
 logger = logging.getLogger(__name__)
 
 
 class TUFDownloader:
-    def __init__(self, verbose=0):
+    def __init__(self, repository_url_prefix=REPOSITORY_URL_PREFIX, verbose=0):
         # 0 => 60 (effectively /dev/null)
         # 1 => 50 (CRITICAL)
         # 2 => 40 (ERROR)
@@ -79,7 +80,17 @@ def __init__(self, verbose=0):
         # https://github.com/theupdateframework/tuf/blob/aa2ab218f22d8682e03c992ea98f88efd155cffd/tuf/client/updater.py#L628-L683
         # NOTE: This updater will store files under:
         # os.path.join(REPOSITORIES_DIR, REPOSITORY_DIR)
-        self.__updater = Updater(REPOSITORY_DIR, REPOSITORY_MIRRORS)
+        self.__updater = Updater(
+            REPOSITORY_DIR,
+            {
+                'repo': {
+                    'url_prefix': repository_url_prefix,
+                    'metadata_path': 'metadata.staged',
+                    'targets_path': 'targets',
+                    'confined_target_dirs': [''],
+                }
+            },
+        )
 
         # NOTE: Update to the latest top-level role metadata only ONCE, so that
         # we use the same consistent snapshot to download targets.

datadog_checks_downloader/tests/test_downloader.py

@@ -8,12 +8,16 @@
 import subprocess
 
 import requests
+import six
+
+from datadog_checks.downloader.download import REPOSITORY_URL_PREFIX
 
 log = logging.getLogger('test_downloader')
 
 
 def test_downloader():
-    r = requests.get('https://dd-integrations-core-wheels-build-stable.datadoghq.com/targets/simple/index.html')
+    index = six.moves.urllib_parse.urljoin(REPOSITORY_URL_PREFIX, 'targets/simple/index.html')
+    r = requests.get(index)
     r.raise_for_status()
 
     for line in r.text.split('\n'):