Refactor CLI code (#3317)

datadog_checks_downloader/datadog_checks/downloader/cli.py

@@ -10,54 +10,18 @@
 # 2nd party.
 from .download import TUFDownloader
 from .exceptions import (
-    InconsistentSimpleIndex,
-    MissingVersions,
     NonCanonicalVersion,
     NonDatadogPackage,
-    NoSuchDatadogPackage,
     NoSuchDatadogPackageOrVersion,
 )
 
 # 3rd party.
-# NOTE: We assume that setuptools is installed by default.
-from pkg_resources import parse_version
 from tuf.exceptions import UnknownTargetError
 
 
 # Private module functions.
 
 
-def __get_latest_version(tuf_downloader, standard_distribution_name, wheel_distribution_name):
-    target_relpath = 'simple/{}/index.html'.format(standard_distribution_name)
-
-    try:
-        # NOTE: We do not perform in-toto inspection for simple indices; only for wheels.
-        target_abspath = tuf_downloader.download(target_relpath, download_in_toto_metadata=False)
-    except UnknownTargetError:
-        raise NoSuchDatadogPackage(standard_distribution_name)
-
-    pattern = "<a href='(" + wheel_distribution_name + "-(.*?)-py2\\.py3-none-any\\.whl)'>(.*?)</a><br />"
-    versions = []
-
-    with open(target_abspath) as simple_index:
-        for line in simple_index:
-            match = re.match(pattern, line)
-            if match:
-                href = match.group(1)
-                version = match.group(2)
-                text = match.group(3)
-                if href != text:
-                    raise InconsistentSimpleIndex(href, text)
-                else:
-                    # https://setuptools.readthedocs.io/en/latest/pkg_resources.html#parsing-utilities
-                    versions.append(parse_version(version))
-
-    if not len(versions):
-        raise MissingVersions(standard_distribution_name)
-    else:
-        return max(versions)
-
-
 def __is_canonical(version):
     '''
     https://www.python.org/dev/peps/pep-0440/#appendix-b-parsing-version-strings-with-regular-expressions
@@ -67,11 +31,14 @@ def __is_canonical(version):
     return re.match(P, version) is not None
 
 
-def __wheel_distribution_name(standard_distribution_name):
+def __get_wheel_distribution_name(standard_distribution_name):
     # https://www.python.org/dev/peps/pep-0491/#escaping-and-unicode
     return re.sub('[^\\w\\d.]+', '_', standard_distribution_name, re.UNICODE)
 
 
+# Public module functions.
+
+
 def download():
     parser = argparse.ArgumentParser()
 
@@ -92,11 +59,11 @@ def download():
     if not standard_distribution_name.startswith('datadog-'):
         raise NonDatadogPackage(standard_distribution_name)
     else:
-        wheel_distribution_name = __wheel_distribution_name(standard_distribution_name)
+        wheel_distribution_name = __get_wheel_distribution_name(standard_distribution_name)
         tuf_downloader = TUFDownloader(verbose=verbose)
 
         if not version:
-            version = __get_latest_version(tuf_downloader, standard_distribution_name, wheel_distribution_name)
+            version = tuf_downloader.get_latest_version(standard_distribution_name, wheel_distribution_name)
         else:
             if not __is_canonical(version):
                 raise NonCanonicalVersion(version)

datadog_checks_downloader/datadog_checks/downloader/download.py

@@ -7,6 +7,7 @@
 import logging
 import logging.config
 import os
+import re
 import shutil
 import tempfile
 
@@ -18,6 +19,7 @@
 
 # Import what we need from TUF.
 from tuf.client.updater import Updater
+from tuf.exceptions import UnknownTargetError
 
 # Import what we need from in-toto.
 from in_toto import verifylib
@@ -30,6 +32,10 @@
     'version': 1,
 })
 
+# Other 3rd-party imports.
+# NOTE: We assume that setuptools is installed by default.
+from pkg_resources import parse_version
+
 # NOTE: A module with a function that substitutes parameters for
 # in-toto inspections. The function is expected to be called
 # 'substitute', and takes one parameter, target_relpath, that specifies
@@ -42,8 +48,11 @@
 
 # Exceptions.
 from .exceptions import (
+    InconsistentSimpleIndex,
+    MissingVersions,
     NoInTotoLinkMetadataFound,
     NoInTotoRootLayoutPublicKeysFound,
+    NoSuchDatadogPackage,
 )
 
 
@@ -163,33 +172,34 @@ def __update_in_toto_layout_pubkeys(self):
         return target_relpaths
 
 
-    def __verify_in_toto_metadata(self, target_relpath, in_toto_metadata_relpaths, pubkey_relpaths):
+    def __verify_in_toto_metadata(self, target_relpath, in_toto_inspection_packet):
         # Make a temporary directory in a parent directory we control.
         tempdir = tempfile.mkdtemp(dir=REPOSITORIES_DIR)
 
+        # Copy files over into temp dir.
+        for rel_path in in_toto_inspection_packet:
+            # Don't confuse Python with any leading path separator.
+            rel_path = rel_path.lstrip('/')
+            abs_path = os.path.join(self.__targets_dir, rel_path)
+            shutil.copy(abs_path, tempdir)
+
+        # Switch to the temp dir.
+        os.chdir(tempdir)
+
+        # Load the root layout and public keys.
+        layout = Metablock.load('root.layout')
+        pubkeys = glob.glob('*.pub')
+        layout_key_dict = import_public_keys_from_files_as_dict(pubkeys)
+        # Parameter substitution.
+        params = substitute(target_relpath)
+
         try:
-            # Copy files over into temp dir.
-            rel_paths = [target_relpath] + in_toto_metadata_relpaths + pubkey_relpaths
-            for rel_path in rel_paths:
-                # Don't confuse Python with any leading path separator.
-                rel_path = rel_path.lstrip('/')
-                abs_path = os.path.join(self.__targets_dir, rel_path)
-                shutil.copy(abs_path, tempdir)
-
-            # Switch to the temp dir.
-            os.chdir(tempdir)
-
-            # Load the root layout and public keys.
-            layout = Metablock.load('root.layout')
-            pubkeys = glob.glob('*.pub')
-            layout_key_dict = import_public_keys_from_files_as_dict(pubkeys)
-            # Verify and inspect.
-            params = substitute(target_relpath)
             verifylib.in_toto_verify(layout, layout_key_dict, substitution_parameters=params)
-            logger.info('in-toto verified {}'.format(target_relpath))
         except:
             logger.exception('in-toto failed to verify {}'.format(target_relpath))
             raise
+        else:
+            logger.info('in-toto verified {}'.format(target_relpath))
         finally:
             # Switch back to a parent directory we control, so that we can
             # safely delete temp dir.
@@ -211,7 +221,11 @@ def __download_and_verify_in_toto_metadata(self, target, target_relpath):
                 raise NoInTotoRootLayoutPublicKeysFound(target_relpath)
 
             else:
-                self.__verify_in_toto_metadata(target_relpath, in_toto_metadata_relpaths, pubkey_relpaths)
+                # Everything we need for in-toto inspection to work: the wheel,
+                # the in-toto root layout, in-toto links, and public keys to
+                # verify the in-toto layout.
+                in_toto_inspection_packet = [target_relpath] + in_toto_metadata_relpaths + pubkey_relpaths
+                self.__verify_in_toto_metadata(target_relpath, in_toto_inspection_packet)
 
 
     def __get_target(self, target_relpath, download_in_toto_metadata=True):
@@ -252,3 +266,39 @@ def download(self, target_relpath, download_in_toto_metadata=True):
             return the complete filepath to the desired target.
         '''
         return self.__get_target(target_relpath, download_in_toto_metadata=download_in_toto_metadata)
+
+
+    def get_latest_version(self, standard_distribution_name, wheel_distribution_name):
+        '''
+        Returns:
+            If download over TUF is successful, this function will return the
+            latest known version of the Datadog integration.
+        '''
+        target_relpath = 'simple/{}/index.html'.format(standard_distribution_name)
+
+        try:
+            # NOTE: We do not perform in-toto inspection for simple indices; only for wheels.
+            target_abspath = self.download(target_relpath, download_in_toto_metadata=False)
+        except UnknownTargetError:
+            raise NoSuchDatadogPackage(standard_distribution_name)
+
+        pattern = "<a href='(" + wheel_distribution_name + "-(.*?)-py2\\.py3-none-any\\.whl)'>(.*?)</a><br />"
+        versions = []
+
+        with open(target_abspath) as simple_index:
+            for line in simple_index:
+                match = re.match(pattern, line)
+                if match:
+                    href = match.group(1)
+                    version = match.group(2)
+                    text = match.group(3)
+                    if href != text:
+                        raise InconsistentSimpleIndex(href, text)
+                    else:
+                        # https://setuptools.readthedocs.io/en/latest/pkg_resources.html#parsing-utilities
+                        versions.append(parse_version(version))
+
+        if not len(versions):
+            raise MissingVersions(standard_distribution_name)
+        else:
+            return max(versions)

datadog_checks_downloader/tests/test_downloader.py

@@ -17,7 +17,7 @@ def test_downloader():
     r.raise_for_status()
 
     for line in r.text.split('\n'):
-        pattern = "<a href='(datadog-\\w+?)/'>\\w+?</a><br />"
+        pattern = r"<a href='(datadog-[\w-]+?)/'>\w+?</a><br />"
         match = re.match(pattern, line)
 
         if match: