add arguments to set proxy and ca_bundle for requests

conda_mirror/conda_mirror.py

@@ -191,6 +191,18 @@ def _make_arg_parser():
         type=int,
         default=1000,
     )
+    ap.add_argument(
+        '--proxy',
+        help=('Proxy URL to access internet if needed'),
+        type=str,
+        default=None,
+    )
+    ap.add_argument(
+        '--ssl_verify',
+        help=('Path to a CA_BUNDLE file with certificates of trusted CAs'),
+        type=str,
+        default=None,
+    )
     return ap
 
 
@@ -270,6 +282,21 @@ def pdb_hook(exctype, value, traceback):
             pdb.post_mortem(traceback)
         sys.excepthook = pdb_hook
 
+    proxies = args.proxy
+    if proxies is not None:
+        # use scheme of proxy url to generate dictionary
+        # if no extra scheme is given
+        # examples:
+        # "http:https://user:pass@proxy.tld"
+        #  -> {'http': 'https://user:pass@proxy.tld'}
+        # "https://user:pass@proxy.tld"
+        #  -> {'https': 'https://user:pass@proxy.tld'}
+        scheme, *url = proxies.split(':')
+        if len(url) > 1:
+            url = ':'.join(url)
+        else:
+            url = '{}:{}'.format(scheme, url[0])
+        proxies = {scheme: url}
     return {
         'upstream_channel': args.upstream_channel,
         'target_directory': args.target_directory,
@@ -281,6 +308,8 @@ def pdb_hook(exctype, value, traceback):
         'dry_run': args.dry_run,
         'no_validate_target': args.no_validate_target,
         'minimum_free_space': args.minimum_free_space,
+        'proxies': proxies,
+        'ssl_verify': args.ssl_verify,
     }
 
 
@@ -361,7 +390,7 @@ def _validate(filename, md5=None, size=None):
     return filename, None
 
 
-def get_repodata(channel, platform):
+def get_repodata(channel, platform, proxies=None, ssl_verify=None):
     """Get the repodata.json file for a channel/platform combo on anaconda.org
 
     Parameters
@@ -370,6 +399,10 @@ def get_repodata(channel, platform):
         anaconda.org/CHANNEL
     platform : {'linux-64', 'linux-32', 'osx-64', 'win-32', 'win-64'}
         The platform of interest
+    proxies : dict
+        Proxys for connecting internet
+    ssl_verify : str
+        Path to a CA_BUNDLE file or directory with certificates of trusted CAs
 
     Returns
     -------
@@ -381,7 +414,7 @@ def get_repodata(channel, platform):
     url = url_template.format(channel=channel, platform=platform,
                               file_name='repodata.json')
 
-    resp = requests.get(url).json()
+    resp = requests.get(url, proxies=proxies, verify=ssl_verify).json()
     info = resp.get('info', {})
     packages = resp.get('packages', {})
     # Patch the repodata.json so that all package info dicts contain a "subdir"
@@ -393,7 +426,7 @@ def get_repodata(channel, platform):
     return info, packages
 
 
-def _download(url, target_directory):
+def _download(url, target_directory, proxies=None, ssl_verify=None):
     """Download `url` to `target_directory`
 
     Parameters
@@ -402,6 +435,10 @@ def _download(url, target_directory):
         The url to download
     target_directory : str
         The path to a directory where `url` should be downloaded
+    proxies : dict
+        Proxys for connecting internet
+    ssl_verify : str
+        Path to a CA_BUNDLE file or directory with certificates of trusted CAs
 
     Returns
     -------
@@ -416,7 +453,8 @@ def _download(url, target_directory):
     download_filename = os.path.join(target_directory, target_filename)
     logger.debug('downloading to %s', download_filename)
     with open(download_filename, 'w+b') as tf:
-        ret = requests.get(url, stream=True)
+        ret = requests.get(url, stream=True,
+                           proxies=proxies, verify=ssl_verify)
         for data in ret.iter_content(chunk_size):
             tf.write(data)
         file_size = os.path.getsize(download_filename)
@@ -547,7 +585,8 @@ def _validate_or_remove_package(args):
 
 def main(upstream_channel, target_directory, temp_directory, platform,
          blacklist=None, whitelist=None, num_threads=1, dry_run=False,
-         no_validate_target=False, minimum_free_space=0):
+         no_validate_target=False, minimum_free_space=0, proxies=None,
+         ssl_verify=None):
     """
 
     Parameters
@@ -589,6 +628,10 @@ def main(upstream_channel, target_directory, temp_directory, platform,
         If True, skip validation of files already present in target_directory.
     minimum_free_space : int, optional
         Stop downloading when free space target_directory or temp_directory reach this threshold.
+    proxies : dict
+        Proxys for connecting internet
+    ssl_verify : str
+        Path to a CA_BUNDLE file or directory with certificates of trusted CAs
 
     Returns
     -------
@@ -649,7 +692,8 @@ def main(upstream_channel, target_directory, temp_directory, platform,
     if not os.path.exists(os.path.join(target_directory, platform)):
         os.makedirs(os.path.join(target_directory, platform))
 
-    info, packages = get_repodata(upstream_channel, platform)
+    info, packages = get_repodata(upstream_channel, platform,
+                                  proxies=proxies, ssl_verify=ssl_verify)
     local_directory = os.path.join(target_directory, platform)
 
     # 1. validate local repo
@@ -739,7 +783,8 @@ def main(upstream_channel, target_directory, temp_directory, platform,
                     break
 
                 # download package
-                total_bytes += _download(url, download_dir)
+                total_bytes += _download(url, download_dir, proxies=proxies,
+                                         ssl_verify=ssl_verify)
 
                 # make sure we have enough free disk space in the target folder to meet threshold
                 # while also being able to fit the packages we have already downloaded