uhd_images_downloader: Add environment variable for http auth

This allows the image downloader to download files from restricted sources
using HTTP basic auth, specifying the credentials in the UHD_IMAGES_USER and
UHD_IMAGES_PASSWORD environment variables:
```
UHD_IMAGES_USER=lane UHD_IMAGES_PASSWORD=MyS3cretPassword uhd_images_downloader.py
```

4 files changed, 79 insertions, 3 deletions

diff --git a/host/cmake/Modules/UHDUnitTest.cmake b/host/cmake/Modules/UHDUnitTest.cmake
index 340279b9b..d477dc966 100644
--- a/host/cmake/Modules/UHDUnitTest.cmake
+++ b/host/cmake/Modules/UHDUnitTest.cmake
@@ -130,6 +130,7 @@ function(UHD_ADD_PYTEST test_name)
             WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/python"
         )
     endif(ENABLE_QEMU_UNITTESTS)
+    # Include ${CMAKE_BINARY_DIR}/utils/ for testing the python utils
     set_tests_properties(${test_name} PROPERTIES
-        ENVIRONMENT PYTHONPATH=${CMAKE_SOURCE_DIR}/tests/common)
+        ENVIRONMENT PYTHONPATH=${CMAKE_SOURCE_DIR}/tests/common:${CMAKE_BINARY_DIR}/utils/)
 endfunction(UHD_ADD_PYTEST)
diff --git a/host/tests/CMakeLists.txt b/host/tests/CMakeLists.txt
index 19d6bd23c..99f4e5eaf 100644
--- a/host/tests/CMakeLists.txt
+++ b/host/tests/CMakeLists.txt
@@ -72,6 +72,7 @@ set(pytest_sources
     pyranges_test.py
     verify_fbs_test.py
     pychdr_parse_test.py
+    uhd_image_downloader_test.py
 )
 
 #turn each test cpp file into an executable with an int main() function
diff --git a/host/tests/uhd_image_downloader_test.py b/host/tests/uhd_image_downloader_test.py
new file mode 100644
index 000000000..ef5c0b5a3
--- /dev/null
+++ b/host/tests/uhd_image_downloader_test.py
@@ -0,0 +1,45 @@
+#
+# Copyright 2020 Ettus Research, a National Instruments Brand
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+#
+"""
+Unit test for uhd_images_downloader
+"""
+
+import unittest
+import uhd_images_downloader
+
+def mk_env(username, password):
+    env = {}
+    if username is not None:
+        env[uhd_images_downloader._USERNAME_VARIABLE] = username
+    if password is not None:
+        env[uhd_images_downloader._PASSWORD_VARIABLE] = password
+    return env
+
+class PyImageDownloaderTest(unittest.TestCase):
+    """ Test Python image downloader functionality """
+    def test_parse_auth_works(self):
+        self.assertEqual(
+            uhd_images_downloader.parse_auth(mk_env("bbrother", "2+2=5")),
+            ("bbrother", "2+2=5")
+        )
+
+    def test_parse_auth_empty_on_empty(self):
+        self.assertEqual(
+            uhd_images_downloader.parse_auth(mk_env(None, None)),
+            None
+        )
+        self.assertEqual(
+            uhd_images_downloader.parse_auth(mk_env("", None)),
+            None
+        )
+
+    def test_throws_on_empty_pw(self):
+        with self.assertRaises(RuntimeError):
+            uhd_images_downloader.parse_auth(mk_env("username", None))
+
+    def test_throws_on_empty_username(self):
+        with self.assertRaises(RuntimeError):
+            uhd_images_downloader.parse_auth(mk_env(None, "password"))
diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index 927d6a481..b9bc2dd0b 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -60,6 +60,8 @@ except ImportError:
     exit(0)
 
 # pylint: disable=bad-whitespace
+_USERNAME_VARIABLE        = "UHD_IMAGES_USER"
+_PASSWORD_VARIABLE        = "UHD_IMAGES_PASSWORD"
 _DEFAULT_TARGET_REGEX     = "(fpga|fw|windrv)_default"
 _BASE_DIR_STRUCTURE_PARTS = ["share", "uhd", "images"]
 _DEFAULT_INSTALL_PATH     = os.path.join("@CMAKE_INSTALL_PREFIX@", *_BASE_DIR_STRUCTURE_PARTS)
@@ -344,6 +346,31 @@ def print_target_list(manifest, args):
             print(args.base_url+manifest_item[1]["url"])
 
 
+def parse_auth(env):
+    """
+    Parse a (potentially None) environment variable string as a username:password
+    for consumption by requests.
+    """
+    username = env.get(_USERNAME_VARIABLE, None)
+    password = env.get(_PASSWORD_VARIABLE, None)
+    if not username and not password:
+        return None
+
+    if not username:
+        raise RuntimeError(
+            "Password variable {} was specified, so username variable {} must be set as well"
+            .format(_PASSWORD_VARIABLE, _USERNAME_VARIABLE)
+        )
+
+    if not password:
+        raise RuntimeError(
+            "Username variable {} was specified, so password variable {} must as well"
+            .format(_USERNAME_VARIABLE, _PASSWORD_VARIABLE)
+        )
+
+    return (username, password)
+
+
 def download(
         images_url,
         filename,
@@ -354,15 +381,17 @@ def download(
     """ Run the download, show progress """
     download_limit = download_limit or _DEFAULT_DOWNLOAD_LIMIT
     log("TRACE", "Downloading {} to {}".format(images_url, filename))
+    auth = parse_auth(os.environ)
     try:
         resp = requests.get(images_url, stream=True, proxies=_PROXIES,
-                            headers={'User-Agent': 'UHD Images Downloader'})
+                            headers={'User-Agent': 'UHD Images Downloader'},
+                            auth=auth)
     except TypeError:
         # requests library versions pre-4c3b9df6091b65d8c72763222bd5fdefb7231149
         # (Dec.'12) workaround
         resp = requests.get(images_url, prefetch=False, proxies=_PROXIES,
                             headers={'User-Agent': 'UHD Images Downloader'},
-                            allow_redirects=True)
+                            allow_redirects=True, auth=auth)
     if resp.status_code != 200:
         raise RuntimeError("URL does not exist: {}".format(images_url))
     filesize = float(resp.headers.get('content-length', -1))