utils: uhd_images_downloader: Add download limit

This download limit prevents users from accidentally downloading really
large files. The default limit is 100 MiB. When a file is selected for
download that exceeds this size, the user is prompted if she really
wants to download that file.

By specifying --yes (or -y), this can can be bypassed. Example:

    $ uhd_images_downloader --yes -t sdimg

Will match all files that match 'sdimg', which are typically larger than
100 MiB, without interaction.

1 files changed, 25 insertions, 1 deletions

diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index 6655274dd..11f8ad39b 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -37,6 +37,7 @@ _DEFAULT_BASE_URL         = "http://files.ettus.com/binaries/cache/"
 _INVENTORY_FILENAME       = "inventory.json"
 _CONTACT                  = "support@ettus.com"
 _DEFAULT_BUFFER_SIZE      = 8192
+_DEFAULT_DOWNLOAD_LIMIT   = 100 * 1024 * 1024 # Bytes
 _ARCHIVE_ALGS             = ["zip", "targz", "tarxz"]
 _ARCHIVE_DEFAULT_TYPE     = "zip"
 _UHD_VERSION              = "@UHD_VERSION@"
@@ -77,6 +78,10 @@ def parse_args():
                         help="With -l, only print the URLs, nothing else.")
     parser.add_argument("--buffer-size", type=int, default=_DEFAULT_BUFFER_SIZE,
                         help="Set download buffer size")
+    parser.add_argument("--download-limit", type=int, default=_DEFAULT_DOWNLOAD_LIMIT,
+                        help="Set threshold for download limits. Any download "
+                             "larger than this will require approval, either "
+                             "interactively, or by providing --yes.")
     parser.add_argument("-b", "--base-url", type=str, default=_DEFAULT_BASE_URL,
                         help="Set base URL for images download location")
     parser.add_argument("-z", "--archive-type", type=str, default=_ARCHIVE_DEFAULT_TYPE,
@@ -220,8 +225,15 @@ def lookup_urls(regex_l, manifest, inventory, refetch=False):
     return selected_targets
 
 
-def download(images_url, filename, buffer_size=_DEFAULT_BUFFER_SIZE, print_progress=False):
+def download(
+        images_url,
+        filename,
+        buffer_size=_DEFAULT_BUFFER_SIZE,
+        print_progress=False,
+        download_limit=None
+    ):
     """ Run the download, show progress """
+    download_limit = download_limit or _DEFAULT_DOWNLOAD_LIMIT
     log("TRACE", "Downloading {} to {}".format(images_url, filename))
     try:
         resp = requests.get(images_url, stream=True,
@@ -234,6 +246,12 @@ def download(images_url, filename, buffer_size=_DEFAULT_BUFFER_SIZE, print_progr
     if resp.status_code != 200:
         raise RuntimeError("URL does not exist: {}".format(images_url))
     filesize = float(resp.headers['content-length'])
+    if filesize > download_limit:
+        if not ask_permission(
+                "The file size for this target ({:.1f} MiB) exceeds the "
+                "download limit ({:.1f} MiB). Continue downloading?".format(
+                    filesize/1024**2, download_limit/1024**2)):
+            return 0, 0, ""
     filesize_dl = 0
     base_filename = os.path.basename(filename)
     if print_progress and not sys.stdout.isatty():
@@ -322,6 +340,9 @@ def main():
     args = parse_args()
     if not args.base_url.endswith('/') and args.base_url != "":
         args.base_url += '/'
+    if args.yes:
+        global _YES
+        _YES = True
     archive_type = args.archive_type
     if archive_type not in _ARCHIVE_ALGS:
         log("ERROR", "Selected archive type not supported: {}".format(archive_type))
@@ -421,6 +442,9 @@ def main():
                         buffer_size=args.buffer_size,
                         print_progress=(_LOG_LEVEL <= _LOG_LEVELS.get("INFO", 3))
                     )
+                    if downloaded_size == 0:
+                        log("INFO", "Skipping target: {}".format(target_name))
+                        continue
                     log("TRACE", "{} successfully downloaded ({} Bytes)"
                         .format(temp_path, downloaded_size))