utils: images downloader: adding SHA256 support

Adding SHA256 checking for downloaded image archives. The expected SHA
is read from the manifest, and checked after the download completes.

1 files changed, 15 insertions, 6 deletions

diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index a64330e15..fcba16c33 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -9,6 +9,7 @@ Download image files required for USRPs
 """
 from __future__ import print_function
 import argparse
+import hashlib
 import json
 import math
 import os
@@ -115,10 +116,10 @@ def parse_manifest(manifest_contents):
             if not line_unpacked or line.strip().startswith('#'):
                 continue
 
-            target, repo_hash, url, sha_hash = line_unpacked
+            target, repo_hash, url, sha256_hash = line_unpacked
             manifest[target] = {"repo_hash": repo_hash,
                                 "url": url,
-                                "sha_hash": sha_hash,
+                                "sha256_hash": sha256_hash,
                                 }
         except ValueError:
             log("WARN", "Warning: Invalid line in manifest file:\n"
@@ -166,6 +167,7 @@ def lookup_urls(regex, manifest, inventory, refetch=False):
             target_info = manifest.get(target)
             target_url = target_info.get("url")
             target_hash = target_info.get("repo_hash")
+            target_sha256 = target_info.get("sha256_hash")
             # Check if the same filename and hash appear in the inventory
             # TODO: Change to the TARGET instead of the filename
             filename = os.path.basename(target_url)
@@ -174,7 +176,9 @@ def lookup_urls(regex, manifest, inventory, refetch=False):
                 log("TRACE", "Selected target already downloaded- skipping.")
             else:
                 # We don't have that exact file, add it to the list
-                selected_targets.append({"repo_hash": target_hash, "url": target_url})
+                selected_targets.append({"repo_hash": target_hash,
+                                         "url": target_url,
+                                         "sha256_hash": target_sha256})
     return selected_targets
 
 
@@ -194,11 +198,13 @@ def download(images_url, filename, buffer_size=_DEFAULT_BUFFER_SIZE, print_progr
     filesize = float(resp.headers['content-length'])
     filesize_dl = 0
     with open(filename, "wb") as temp_file:
+        sha256_sum = hashlib.sha256()
         base_filename = os.path.basename(filename)
         for buff in resp.iter_content(chunk_size=buffer_size):
             if buff:
                 temp_file.write(buff)
                 filesize_dl += len(buff)
+                sha256_sum.update(buff)
             if print_progress:
                 status = r"%05d kB / %05d kB (%03d%%) %s" % (
                     int(math.ceil(filesize_dl / 1000.)), int(math.ceil(filesize / 1000.)),
@@ -212,7 +218,7 @@ def download(images_url, filename, buffer_size=_DEFAULT_BUFFER_SIZE, print_progr
                 sys.stdout.flush()
     if print_progress:
         print('')
-    return filesize, filesize_dl
+    return filesize, filesize_dl, sha256_sum.hexdigest()
 
 
 def delete_from_inv(archive_fn, inventory, images_dir):
@@ -330,18 +336,21 @@ def main():
             for target_info in targets_info:
                 target_hash = target_info.get("repo_hash")
                 target_rel_url = target_info.get("url")
+                target_sha256 = target_info.get("sha256_hash")
                 filename = os.path.basename(target_rel_url)
                 temp_path = os.path.join(temp_dir, filename)
                 # Add a trailing slash to make sure that urljoin handles things properly
                 full_url = urljoin(args.base_url+'/', target_rel_url)
                 if not args.dry_run:
-                    _, downloaded_size = download(
+                    _, downloaded_size, downloaded_sha256 = download(
                         images_url=full_url,
                         filename=temp_path,
                         buffer_size=args.buffer_size,
                         print_progress=(_LOG_LEVEL <= _LOG_LEVELS.get("DEBUG", 2))
                     )
-                    # TODO: Check SHA
+                    if downloaded_sha256 != target_sha256:
+                        log("ERROR", "Downloaded SHA256 does not match manifest for {}!".format(
+                            full_url))
                     log("TRACE", "{} successfully downloaded ({} Bytes)"
                         .format(temp_path, downloaded_size))