kbenv: allow additional checksum verification types

The k8s release process has changed the provided checksum types a few
times. Currently only sha256 & sha512 files are provided, but in the
recent past only md5 & sha1 files were provided.

This new verification process should account for both old and new
kubectl versions.

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ### Fixed
 - fixed `TypeError` when Static Site is configured to use the S3 bucket to serve the website directly (e.g. CloudFront disabled)
+- Newer kubectl versions download errors (now supporting sha1/256/512 checksum verification)
 
 ## [1.16.0] - 2020-11-09
 ### Added

runway/env_mgr/kbenv.py

@@ -6,20 +6,96 @@
 import sys
 import tempfile
 
+import requests
+
 # Old pylint on py2.7 incorrectly flags these
 from six.moves.urllib.error import URLError  # pylint: disable=E
 from six.moves.urllib.request import urlretrieve  # pylint: disable=E
 
-from ..util import md5sum
+from ..util import get_file_hash
 from . import EnvManager, handle_bin_download_error
 
 LOGGER = logging.getLogger(__name__)
 KB_VERSION_FILENAME = ".kubectl-version"
 RELEASE_URI = "https://storage.googleapis.com/kubernetes-release/release"
 
 
-# Branch and local variable count will go down when py2 support is dropped
-def download_kb_release(  # noqa pylint: disable=too-many-locals,too-many-branches
+def verify_kb_release(kb_url, download_dir, filename):
+    """Compare checksum and exit if it doesn't match.
+
+    Different releases provide varying checksum files. To account for this,
+    start at SHA512 and work down to the first available checksum.
+
+    requests is used for downloading these small files because of difficulty in
+    getting 404 status from urllib on py2. Once py2 support is dropped, downloads
+    can be moved to urllib.
+
+    https://stackoverflow.com/questions/1308542/how-to-catch-404-error-in-urllib-urlretrieve
+
+    """
+    # This might be a bit cleaner refactored as self-referencing function, but
+    # the ridiculousness should be short-lived as md5 & sha1 support won't last
+    # long.
+    try:
+        checksum_type = "sha512"
+        checksum_filename = filename + "." + checksum_type
+        LOGGER.debug("attempting download of kubectl %s checksum...", checksum_type)
+        download_request = requests.get(
+            kb_url + "/" + checksum_filename, allow_redirects=True
+        )
+        download_request.raise_for_status()
+    except requests.exceptions.HTTPError:
+        try:
+            checksum_type = "sha256"
+            checksum_filename = filename + "." + checksum_type
+            LOGGER.debug("attempting download of kubectl %s checksum...", checksum_type)
+            download_request = requests.get(
+                kb_url + "/" + checksum_filename, allow_redirects=True
+            )
+            download_request.raise_for_status()
+        except requests.exceptions.HTTPError:
+            try:
+                checksum_type = "sha1"
+                checksum_filename = filename + "." + checksum_type
+                LOGGER.debug(
+                    "attempting download of kubectl %s checksum...", checksum_type
+                )
+                download_request = requests.get(
+                    kb_url + "/" + checksum_filename, allow_redirects=True
+                )
+                download_request.raise_for_status()
+            except requests.exceptions.HTTPError:
+                try:
+                    checksum_type = "md5"
+                    checksum_filename = filename + "." + checksum_type
+                    LOGGER.debug(
+                        "attempting download of kubectl %s checksum...", checksum_type
+                    )
+                    download_request = requests.get(
+                        kb_url + "/" + checksum_filename, allow_redirects=True
+                    )
+                    download_request.raise_for_status()
+                except requests.exceptions.HTTPError:
+                    LOGGER.error("Unable to retrieve kubectl checksum file")
+                    sys.exit(1)
+
+    if sys.version_info < (3, 0):
+        kb_hash = download_request.content.rstrip("\n")
+    else:
+        kb_hash = download_request.content.decode().rstrip("\n")
+
+    if kb_hash != get_file_hash(os.path.join(download_dir, filename), checksum_type):
+        LOGGER.error(
+            "downloaded kubectl %s does not match %s checksum %s",
+            filename,
+            checksum_type,
+            kb_hash,
+        )
+        sys.exit(1)
+    LOGGER.debug("kubectl matched %s checksum...", checksum_type)
+
+
+def download_kb_release(
     version, versions_dir, kb_platform=None, arch=None,
 ):
     """Download kubectl and return path to it."""
@@ -48,18 +124,12 @@ def download_kb_release(  # noqa pylint: disable=too-many-locals,too-many-branch
 
     try:
         LOGGER.verbose("downloading kubectl from %s...", kb_url)
-        for i in [filename, filename + ".md5"]:
-            urlretrieve(kb_url + "/" + i, os.path.join(download_dir, i))
+        urlretrieve(kb_url + "/" + filename, os.path.join(download_dir, filename))
     # IOError in py2; URLError in 3+
     except (IOError, URLError) as exc:
         handle_bin_download_error(exc, "kubectl")
 
-    with open(os.path.join(download_dir, filename + ".md5"), "r") as stream:
-        kb_hash = stream.read().rstrip("\n")
-
-    if kb_hash != md5sum(os.path.join(download_dir, filename)):
-        LOGGER.error("downloaded kubectl %s does not match md5 %s", filename, kb_hash)
-        sys.exit(1)
+    verify_kb_release(kb_url, download_dir, filename)
 
     version_dir.mkdir(parents=True, exist_ok=True)
     shutil.move(os.path.join(download_dir, filename), str(version_dir / filename))

runway/util.py

@@ -721,6 +721,18 @@ def run_commands(
                     sys.exit(1)
 
 
+def get_file_hash(filename, algorithm):
+    """Return cryptographic hash of file."""
+    file_hash = getattr(hashlib, algorithm)()
+    with open(filename, "rb") as stream:
+        while True:
+            data = stream.read(65536)  # 64kb chunks
+            if not data:
+                break
+            file_hash.update(data)
+    return file_hash.hexdigest()
+
+
 def md5sum(filename):
     """Return MD5 hash of file."""
     md5 = hashlib.md5()