Support multiple root layouts

datadog_checks_downloader/datadog_checks/downloader/cli.py

@@ -8,7 +8,7 @@
 import re
 
 # 2nd party.
-from .download import REPOSITORY_URL_PREFIX, TUFDownloader
+from .download import DEFAULT_ROOT_LAYOUT_TYPE, REPOSITORY_URL_PREFIX, ROOT_LAYOUTS, TUFDownloader
 from .exceptions import NonCanonicalVersion, NonDatadogPackage
 
 # Private module functions.
@@ -39,6 +39,14 @@ def download():
 
     parser.add_argument('--version', type=str, default=None, help='The version number of the desired Datadog check.')
 
+    parser.add_argument(
+        '--type',
+        type=str,
+        default=DEFAULT_ROOT_LAYOUT_TYPE,
+        choices=list(ROOT_LAYOUTS),
+        help='The type of integration.',
+    )
+
     parser.add_argument(
         '-v', '--verbose', action='count', default=0, help='Show verbose information about TUF and in-toto.'
     )
@@ -47,6 +55,7 @@ def download():
     repository_url_prefix = args.repository
     standard_distribution_name = args.standard_distribution_name
     version = args.version
+    root_layout_type = args.type
     verbose = args.verbose
 
     if not standard_distribution_name.startswith('datadog-'):
@@ -55,7 +64,9 @@ def download():
     if version and not __is_canonical(version):
         raise NonCanonicalVersion(version)
 
-    tuf_downloader = TUFDownloader(repository_url_prefix=repository_url_prefix, verbose=verbose)
+    tuf_downloader = TUFDownloader(
+        repository_url_prefix=repository_url_prefix, root_layout_type=root_layout_type, verbose=verbose
+    )
     wheel_relpath = tuf_downloader.get_wheel_relpath(standard_distribution_name, version=version)
     wheel_abspath = tuf_downloader.download(wheel_relpath)
     print(wheel_abspath)  # pylint: disable=print-statement

datadog_checks_downloader/datadog_checks/downloader/download.py

@@ -23,6 +23,7 @@
 from .exceptions import (
     DuplicatePackage,
     InconsistentSimpleIndex,
+    IncorrectRootLayoutType,
     MissingVersions,
     NoInTotoLinkMetadataFound,
     NoInTotoRootLayoutPublicKeysFound,
@@ -48,15 +49,18 @@
 REPOSITORY_URL_PREFIX = 'https://dd-integrations-core-wheels-build-stable.datadoghq.com'
 # Where to find our in-toto root layout.
 IN_TOTO_METADATA_DIR = 'in-toto-metadata'
-IN_TOTO_ROOT_LAYOUT = '2.root.layout'
+ROOT_LAYOUTS = {'core': '2.root.layout', 'extras': '1.extras.root.layout'}
+DEFAULT_ROOT_LAYOUT_TYPE = 'core'
 
 
 # Global variables.
 logger = logging.getLogger(__name__)
 
 
 class TUFDownloader:
-    def __init__(self, repository_url_prefix=REPOSITORY_URL_PREFIX, verbose=0):
+    def __init__(
+        self, repository_url_prefix=REPOSITORY_URL_PREFIX, root_layout_type=DEFAULT_ROOT_LAYOUT_TYPE, verbose=0
+    ):
         # 0 => 60 (effectively /dev/null)
         # 1 => 50 (CRITICAL)
         # 2 => 40 (ERROR)
@@ -71,6 +75,9 @@ def __init__(self, repository_url_prefix=REPOSITORY_URL_PREFIX, verbose=0):
 
         tuf_settings.repositories_directory = REPOSITORIES_DIR
 
+        self.root_layout_type = root_layout_type
+        self.root_layout = ROOT_LAYOUTS[self.root_layout_type]
+
         # NOTE: The directory where the targets for *this* repository is
         # cached. We hard-code this keep this to a subdirectory dedicated to
         # this repository.
@@ -124,7 +131,7 @@ def __download_in_toto_root_layout(self):
         # expected version of the root layout. This is so that, for example, we
         # can introduce new parameters w/o breaking old downloaders that don't
         # know how to substitute them.
-        target_relpath = os.path.join(IN_TOTO_METADATA_DIR, IN_TOTO_ROOT_LAYOUT)
+        target_relpath = os.path.join(IN_TOTO_METADATA_DIR, self.root_layout)
         return self.__download_with_tuf(target_relpath)
 
     def __download_custom(self, target, extension):
@@ -133,6 +140,11 @@ def __download_custom(self, target, extension):
 
         fileinfo = target.get('fileinfo', {})
         custom = fileinfo.get('custom', {})
+
+        root_layout_type = custom.get('root-layout-type', DEFAULT_ROOT_LAYOUT_TYPE)
+        if root_layout_type != self.root_layout_type:
+            raise IncorrectRootLayoutType(root_layout_type, self.root_layout_type)
+
         in_toto_metadata = custom.get('in-toto', [])
 
         for target_relpath in in_toto_metadata:
@@ -177,7 +189,7 @@ def __download_in_toto_links(self, target, target_relpath):
             return link_abspaths
 
     def __load_root_layout(self, target_relpath):
-        root_layout = Metablock.load(IN_TOTO_ROOT_LAYOUT)
+        root_layout = Metablock.load(self.root_layout)
         root_layout_pubkeys = glob.glob('*.pub')
         root_layout_pubkeys = import_public_keys_from_files_as_dict(root_layout_pubkeys)
         # Parameter substitution.
@@ -188,9 +200,9 @@ def __handle_in_toto_verification_exception(self, target_relpath, e):
         logger.exception('in-toto failed to verify %s', target_relpath)
 
         if isinstance(e, LinkNotFoundError) and str(e) == RevokedDeveloper.MSG:
-            raise RevokedDeveloper(target_relpath, IN_TOTO_ROOT_LAYOUT)
+            raise RevokedDeveloper(target_relpath, self.root_layout)
         else:
-            raise
+            raise e
 
     def __in_toto_verify(self, inspection_packet, target_relpath):
         # Make a temporary directory in a parent directory we control.

datadog_checks_downloader/datadog_checks/downloader/exceptions.py

@@ -29,6 +29,15 @@ def __str__(self):
 # Exceptions for the download module.
 
 
+class IncorrectRootLayoutType(Exception):
+    def __init__(self, found, expected):
+        self.found = found
+        self.expected = expected
+
+    def __str__(self):
+        return 'found {}, expected {}'.format(self.found, self.expected)
+
+
 class SimpleIndexError(Exception):
     def __init__(self, standard_distribution_name):
         self.standard_distribution_name = standard_distribution_name