Support ca_cert option for downloads

Author: arjunsuresh

automation/utils.py

@@ -188,40 +188,57 @@ def download_file(i):
     else:
         verify = is_true(i.get('verify_ssl', True))
 
+    # Support custom CA file via MLC_SSL_CA_FILE
+    ca_file = os.environ.get('MLC_SSL_CA_FILE', i.get('ssl_ca_file', ''))
+    if ca_file and os.path.isfile(ca_file):
+        verify = ca_file
+
     try:
-        with requests.get(url, stream=True, allow_redirects=True, verify=verify) as download:
+        try:
+            download = requests.get(url, stream=True, allow_redirects=True, verify=verify)
             download.raise_for_status()
+        except requests.exceptions.SSLError as ssl_err:
+            if verify is not False:
+                print(f"WARNING: SSL verification failed: {ssl_err}")
+                print("Retrying without SSL verification...")
+                download = requests.get(url, stream=True, allow_redirects=True, verify=False)
+                download.raise_for_status()
+            else:
+                raise
 
-            size_string = download.headers.get('Content-Length')
-
-            if size_string is None:
-                transfer_encoding = download.headers.get(
-                    'Transfer-Encoding', '')
-                if transfer_encoding != 'chunked':
-                    return {'return': 1, 'error': 'did not receive file'}
-                else:
-                    size_string = "0"
+        size_string = download.headers.get('Content-Length')
 
-            size = int(size_string)
+        if size_string is None:
+            transfer_encoding = download.headers.get(
+                'Transfer-Encoding', '')
+            if transfer_encoding != 'chunked':
+                return {'return': 1, 'error': 'did not receive file'}
+            else:
+                size_string = "0"
 
-            with open(path_to_file, 'wb') as output:
-                for chunk in download.iter_content(chunk_size=chunk_size):
+        size = int(size_string)
 
-                    if chunk:
-                        output.write(chunk)
-                    if size == 0:
-                        continue
-                    downloaded += 1
-                    percent = downloaded * chunk_size * 100 / size
+        with open(path_to_file, 'wb') as output:
+            for chunk in download.iter_content(chunk_size=chunk_size):
 
-                    sys.stdout.write("\r{}{:3.0f}%".format(text, percent))
-                    sys.stdout.flush()
+                if chunk:
+                    output.write(chunk)
+                if size == 0:
+                    continue
+                downloaded += 1
+                percent = downloaded * chunk_size * 100 / size
 
-                sys.stdout.write("\r{}{:3.0f}%".format(text, 100))
+                sys.stdout.write("\r{}{:3.0f}%".format(text, percent))
                 sys.stdout.flush()
 
+            sys.stdout.write("\r{}{:3.0f}%".format(text, 100))
+            sys.stdout.flush()
+
     except Exception as e:
         return {'return': 1, 'error': format(e)}
+    finally:
+        if 'download' in dir():
+            download.close()
 
     print('')
     if size == 0:

script/download-file/customize.py

@@ -92,6 +92,13 @@ def preprocess(i):
         else:
             verify_ssl = True
 
+        # Support custom CA file via MLC_SSL_CA_FILE
+        ca_file = env.get('MLC_SSL_CA_FILE', '')
+        if ca_file and os.path.isfile(ca_file):
+            ssl_ca_file = ca_file
+        else:
+            ssl_ca_file = ''
+
         if env.get('MLC_DOWNLOAD_PATH', '') != '':
             download_path = env['MLC_DOWNLOAD_PATH']
             if os.path.isfile(download_path):
@@ -186,7 +193,8 @@ def preprocess(i):
                     for i in range(1, 5):
                         r = download_file({
                             'url': url,
-                            'verify': verify_ssl})
+                            'verify': verify_ssl,
+                            'ssl_ca_file': ssl_ca_file})
                         if r['return'] == 0:
                             download_success = True
                             break
@@ -213,7 +221,9 @@ def preprocess(i):
         elif tool == "wget":
             if env.get('MLC_DOWNLOAD_FILENAME', '') != '':
                 extra_download_options += f" --tries=3 -O {q}{env['MLC_DOWNLOAD_FILENAME']}{q} "
-                if not verify_ssl:
+                if ssl_ca_file:
+                    extra_download_options += f" --ca-certificate={q}{ssl_ca_file}{q} "
+                elif not verify_ssl:
                     extra_download_options += " --no-check-certificate "
             env['MLC_DOWNLOAD_CMD'] = f"wget -nc {extra_download_options} {url}"
             for i in range(1, 5):
@@ -238,6 +248,10 @@ def preprocess(i):
             env['MLC_DOWNLOAD_CMD'] += f" -d {q}{os.path.join(os.getcwd(), temp_download_file)}{q} {extra_download_options} {url}"
 
         elif tool == "curl":
+            if ssl_ca_file:
+                extra_download_options += f" --cacert {q}{ssl_ca_file}{q} "
+            elif not verify_ssl:
+                extra_download_options += " -k "
             if env.get('MLC_DOWNLOAD_FILENAME', '') != '':
                 extra_download_options += f" --output {q}{env['MLC_DOWNLOAD_FILENAME']}{q} "