########################################################################
# Copyright (C) 2008-2023 VMware, Inc.
# All Rights Reserved
########################################################################
#

import io
import logging
import os
import shutil
import sys
import time
import zipfile

if sys.version_info[0] >= 3:
   from urllib.request import url2pathname, build_opener, HTTPError
   from urllib.parse import splitport, urlparse
   if os.name == 'nt':
      from urllib.request import proxy_bypass_registry
else:
   from urllib import url2pathname, splitport
   from urllib2 import build_opener, HTTPError
   from urlparse import urlparse
   if os.name == 'nt':
      from urllib import proxy_bypass_registry

from . import Errors
from .Utils import PathUtils, Proxy

try:
   import esxclipy
   esxcliAvailable = True
except ImportError:
   esxcliAvailable = False

CA_CERT_DIR = os.path.join(os.path.sep, 'etc', 'ssl', 'certs')

log = logging.getLogger('downloader')

# Set default options
options = {'timeout': 60,    # seconds
           'retry':   10,
           'throttle': 0}

class DownloaderError(Exception):
   def __init__(self, url, local, err):
      Exception.__init__(self, url, local, err)
      self.url = url
      self.local = local
      self.err = err

class UnzipError(Exception):
   def __init__(self, zipfile, msg):
      Exception.__init__(self, zipfile, msg)
      self.zipfile = zipfile
      self.msg = msg

class _ThrottledDownload(object):
   '''Download object class that implement work-and-wait throttling'''
   def __init__(self, robj, throttle):
      self.robj = robj
      self.throttle = throttle
      self.start = time.time()
      self.downloaded = 0

   def _read_and_delay(self, length):
      '''Download a chunk and delay according to throttle parameter'''
      data = self.robj.read(length)
      elapsed_time = time.time() - self.start
      expected_time = self.downloaded / self.throttle
      if elapsed_time < expected_time:
         time.sleep(expected_time - elapsed_time)
      return data

   def read(self, length=None):
      if length == 0:
         return b''
      KB = 1024
      start = self.downloaded
      downloaded = 0 # record this read length
      data = b''

      while True:
         # read size can be standard or what is left
         if length:
            read_size = min(10 * KB, length - downloaded)
         else:
            read_size = 10 * KB
         newdata = self._read_and_delay(read_size)
         downloaded += len(newdata)
         data += newdata
         if not newdata or downloaded == length:
            break
      return data

   def close(self):
      self.robj.close()

   def _unsupported(self, *arg, **kwargs):
      raise io.UnsupportedOperation

   write = seek = truncate = _unsupported

class Downloader(object):
   '''Simple class to download a URL to the local filesystem or open a
      file stream from a URL.  Can update a progress bar.
   '''
   def __init__(self, url, local='', fileobj=None, progress=None, **opts):
      '''* url      - URL to retrieve.
         * local    - Where to save retrieved file.  Set to '' for streaming
                      mode.
         * fileobj  - File descriptor to save file to, usually a temporary
                      file object. Optional when local is provided or in
                      streaming mode. Caller is responsible of closing the
                      descriptor.
         * progress - If provided, should be on object which provides the
                      following methods:
                      * start()    - Called at the beginning of download with
                                     no arguments.
                      * set()      - Takes two arguments, the first is the
                                     amount completed and the second is the
                                     total.
                      * show()     - Takes no arguments, and causes the
                                     progress to be shown on screen.
                      * setTopic() - Takes one argument, the name of the file
                                     being retrieved.
                      * stop()     - Takes no arguments, called at finish.
         * opts     - Extra options to be parsed, currently support timeout,
                      retries, proxies, ssl context.
                      They only affect this instance of Downloader and not
                      the global options of this module.
      '''
      # handle file:/// as regular paths.
      scheme, netloc, path, params, query, fragment = urlparse(url)
      if scheme == "file":
         self.url = url2pathname(path)
      else:
         self.url = url
      host, port = splitport(netloc)
      # deal with IPv6 url case
      host.strip("[]")
      self.netloc = "%s:%s" % (host, port)
      self.local = local
      self.fileobj = fileobj
      self.progress = progress
      # update options
      self.options = options
      self._updateOptions(opts)

   def _updateOptions(self, opts):
      '''Update options for the downloader instance'''
      # filter None out
      opts = dict(filter(lambda item: item[1] is not None, opts.items()))
      # filter invalid options out
      if int(opts.get('timeout', 0)) <= 0:
         opts.pop('timeout', 0)
      if int(opts.get('retry', 1)) <= 0:
         opts.pop('retry', 1)
      if int(opts.get('throttle', 0)) < 0:
         opts.pop('throttle', 0)
      # update options
      self.options.update(opts)

   def _getUrlOpener(self):
      """Return an OpenerDirector which can work around https and use proxy"""
      if os.name == 'nt' and proxy_bypass_registry(self.netloc):
         proxies = {}
      else:
         proxies = self.options.get('proxies')
         if proxies:
            proxiesCensored = dict(proxies)
            for pt in proxiesCensored:
               proxy = proxiesCensored[pt]
               urlObj = urlparse(proxy)
               if urlObj.password:
                  proxiesCensored[pt] = proxy.replace(urlObj.password, '****')
            log.info('Proxies configured: %s', str(proxiesCensored))
      proxy_handler = Proxy.ProxyHandler(proxies)

      if sys.version_info[0] >= 3:
         from urllib.request import HTTPSHandler
      else:
         from urllib2 import HTTPSHandler

      # create https handler with configured ssl context.
      # use unverified ssl context if ssl context is
      # not configured.
      sslcontext = self._getsslcontext()
      if not sslcontext:
         import ssl
         sslcontext = ssl._create_unverified_context()
      https_handler = HTTPSHandler(context=sslcontext)
      return build_opener(proxy_handler, https_handler)

   def _urlopen(self, url):
      """Wrapper to urlopen, to handle HTTPS and proxy.
      """
      opener = self._getUrlOpener()

      # use timeout when available
      try:
         return opener.open(url, timeout=self.options['timeout'])
      except KeyError:
         return opener.open(url)

   def _retry(self, func, retries):
      '''Retry operation'''
      def _checkRetry(attemptNum, maxRetries, exception):
         '''Check retry count and raise when it is the last retry
         '''
         if attemptNum == maxRetries:
            raise exception
         time.sleep(3)
         log.warning('Download failed: %s, %u retry left...', str(exception),
                     maxRetries - attemptNum)
      for retry in range(1, retries + 1):
         try:
            return func()
         except HTTPError as err:
            # Check if we should not retry on a "404: Not Found" error.
            if err.code == 404:
               if self.options.get('disable404Retry', False):
                  log.info('Download failed: file not found (error 404). Retry '
                           'is disabled.')
                  raise
            _checkRetry(retry, retries, err)
         except Exception as e:
            _checkRetry(retry, retries, e)

   def _generate_url_request_with_ssl_context(self, sslcontext,
                                              proxies, certExists):
      '''Generate a get request for the URL with an SSL session
         that honors the SSL context configured for the Downloader.
         This will be used in the context of a vCenter only.
         ESXi doesn't have requests library, so this will throw ImportError.
         Parameters:
            * sslcontext - An object of type ssl.SSLContext.
            * proxies - The proxy configuration that will be used.
            * certExists - Set to True to enfore certificate validation of
                           the server during download.
      '''
      import ssl
      from urllib3.poolmanager import PoolManager
      from urllib3.util import parse_url
      from requests.sessions import Session
      from requests.adapters import HTTPAdapter
      from requests.utils import prepend_scheme_if_needed, select_proxy

      class SslAwareAdapter(HTTPAdapter):
         '''An SSL aware HTTPAdapter that enforces SSL context support
            in requests.
         '''
         def __init__(self, sslContext=None, **kwargs):
            '''Initializes the SSL aware context.
               Parameters:
                  * sslContext - An object of type ssl.SSLContext.
            '''
            self.sslContext = sslContext
            super().__init__(**kwargs)

         def init_poolmanager(self, connections, maxsize, block=False):
            '''Initializes the pool of connections that handle the requests.
               This function overrides the default implementation of
               init_poolmanager in HTTPAdapter to set the configured SSL
               context for the connections.
               Parameters:
                  * connections - The number of urllib3 connection pools to
                                  cache.
                  * maxsize - The maximum number of connections to save in the
                              pool.
                  * block - Block when no free connections are available.
            '''
            self.poolmanager = PoolManager(block=block,
                                           maxsize=maxsize,
                                           num_pools=connections,
                                           ssl_context=self.sslContext)

         def get_connection(self, url, proxies=None):
            '''Returns the thread-safe connection pool.
               This function overrides the default implementation in HTTPAdapter
               to set the configured SSL context if ProxyManager is used instead
               of PoolManager. The default implementation ignores(!) the SSL
               context if the ProxyManager is used.
               Parameters:
                  * url - The URL to connect to.
                  * proxies - A Requests-style dictionary of proxies used on
                              this request.
               Returns:
                  * An object of type urllib3.ConnectionPool
            '''
            proxy = select_proxy(url, proxies)
            # If proxies are not configured, use the default implementation
            # from the library.
            if not proxy:
               return super().get_connection(url, proxies)

            proxy = prepend_scheme_if_needed(proxy, 'http')
            proxy_url = parse_url(proxy)
            # If proxies are configured, create ProxyManager that will honor
            # the configured SSL context. (The library implementation ignores
            # the SSL context).
            proxy_manager = self.proxy_manager_for(proxy,
                                                   ssl_context=self.sslContext)
            return proxy_manager.connection_from_url(url)

      session = Session()
      session.mount("https://", SslAwareAdapter(sslcontext))
      with session:
         if certExists and proxies:
            req = session.request(method='get', url=self.url, stream=True,
                                  allow_redirects=True, proxies=proxies,
                                  verify=CA_CERT_DIR)
         elif proxies:
            req = session.request(method='get', url=self.url, stream=True,
                                  allow_redirects=True, proxies=proxies)
         elif certExists:
            req = session.request(method='get', url=self.url, stream=True,
                                  allow_redirects=True, verify=CA_CERT_DIR)
         else:
            req = session.request(method='get', url=self.url, stream=True,
                                  allow_redirects=True)
         self.url = req.url
         return req

   def _generate_url_request(self):
      '''Generate a get request for the URL.
      '''
      # Import should fail for empty requests folder
      from requests import get, utils
      proxies = self.options.get('proxies', None)
      certExists = os.path.exists(CA_CERT_DIR)
      if os.name == 'nt' and proxy_bypass_registry(self.netloc):
         proxies = None

      # The module "requests" has issue with no_proxy setting.
      # Thus, we workaround the issue here.
      if (proxies and 'no_proxy' in proxies and
          utils.should_bypass_proxies(self.url, proxies['no_proxy'])):
         proxies = None

      sslcontext = self._getsslcontext()
      # If an SSL context was set for Downloader, generate
      # requests using the SSL context.
      if sslcontext:
         try:
            req = self._generate_url_request_with_ssl_context(sslcontext,
                                                              proxies,
                                                              certExists)
            self.url = req.url
            return req
         except Exception as ex:
            raise DownloaderError(self.url, self.local, str(ex))

      if certExists and proxies:
         req = get(self.url, stream=True, allow_redirects=True,
                            proxies=proxies, verify=CA_CERT_DIR)
      elif proxies:
         req = get(self.url, stream=True, allow_redirects=True,
                            proxies=proxies)
      elif certExists:
         req = get(self.url, stream=True, allow_redirects=True,
                            verify=CA_CERT_DIR)
      else:
         req = get(self.url, stream=True, allow_redirects=True)

      self.url = req.url
      return req

   def _open_for_download(self):
      '''Open file object for download'''
      try:
         return self._generate_url_request()
      except ImportError:
         # Module "requests" doesn't exists. Fall back to old code.
         pass

      robj = self._urlopen(self.url)
      self.url = robj.geturl()
      # when throttled, use custom wrapper
      if self.options['throttle'] > 0:
         return _ThrottledDownload(robj, self.options['throttle'])
      else:
         return robj

   def _download_to_file(self):
      '''Download to file'''
      READ_SIZE = 32 * 1024
      def doRequestsDownload(req, fobj):
         """Download from request to fobj with requests module.
         """
         if self.progress:
            self.progress.start()
         downloaded = 0
         for chunk in req.iter_content(READ_SIZE):
            downloaded += len(chunk)
            if self.progress:
               self.progress.set(downloaded, totalSize)
               self.progress.show()
            fobj.write(chunk)
         if self.progress:
            self.progress.stop()
         if totalSize > 0 and downloaded < totalSize:
            raise DownloaderError(self.url, self.local,
               'Incomplete download from %s, received %u bytes, '
               'expecting %u bytes' % (self.url, downloaded, totalSize))

      try:
         with self._generate_url_request() as req:
            totalSize = int(req.headers.get('Content-Length', '0').strip())
            if self.fileobj:
               doRequestsDownload(req, self.fileobj)
               self.fileobj.seek(0)
            else:
               with open(self.local, 'wb') as fobj:
                  doRequestsDownload(req, fobj)
         return
      except ImportError:
         # Module "requests" doesn't exists. Fall back to old code.
         pass

      def doDownload(robj, fobj):
         """Download from robj to fobj.
         """
         if self.progress:
            self.progress.start()
         downloaded = 0
         data = robj.read(READ_SIZE)
         while data:
            downloaded += len(data)
            if self.progress:
               self.progress.set(downloaded, totalSize)
               self.progress.show()
            fobj.write(data)
            data = robj.read(READ_SIZE)
         if self.progress:
            self.progress.stop()
         if totalSize > 0 and downloaded < totalSize:
            raise DownloaderError(self.url, self.local,
               'Incomplete download from %s, received %u bytes, expecting '
               '%u bytes' % (self.url, downloaded, totalSize))

      robj = self._urlopen(self.url)
      self.url = robj.geturl()
      # get length header
      totalSize = int(robj.headers.get('Content-Length', '0').strip())

      # with rate limit, wrap with custom object
      if self.options['throttle'] > 0:
         robj = _ThrottledDownload(robj, self.options['throttle'])

      try:
         if self.fileobj:
            doDownload(robj, self.fileobj)
            self.fileobj.seek(0)
         else:
            with open(self.local, 'wb') as fobj:
               doDownload(robj, fobj)
      finally:
         robj.close()

   def _getfromurl(self):
      # Translate datastore path to absolute file path
      if PathUtils.IsDatastorePath(self.url):
         self.url = PathUtils.DatastoreToFilepath(self.url)

      if os.path.exists(self.url):
         self.local = os.path.normpath(self.url)
         return self.local

      log.info('Downloading %s to %s' % (self.url, self.local))
      try:
         self._retry(self._download_to_file, self.options['retry'])
      except Exception as e:
         raise DownloaderError(self.url, self.local, str(e))
      return self.local

   def _openfromurl(self):
      # Translate datastore path to absolute file path
      if PathUtils.IsDatastorePath(self.url):
         self.url = PathUtils.DatastoreToFilepath(self.url)

      if os.path.exists(self.url):
         return open(self.url, 'rb')

      log.info('Opening %s for download' % self.url)
      try:
         return self._retry(self._open_for_download, self.options['retry'])
      except Exception as e:
         raise DownloaderError(self.url, self.local, str(e))

   @staticmethod
   def _openzipfile(zipfn, memberfn):
      z = zipfile.ZipFile(zipfn)
      namelist = z.namelist()
      if memberfn not in namelist:
         # Try to match ./path/to/meta, path/to/meta, /path/to/meta.
         normalized = memberfn.lstrip(os.sep + os.curdir)
         for name in namelist:
            if name.lstrip(os.sep + os.curdir) == normalized:
               memberfn = name
               break
      # Python 2.6 only
      return z.open(memberfn)

   def _openfromzip(self):
      try:
         zipfn, memberfn = self.url[4:].split('?', 1)
      except Exception:
         msg = "Invalid Location: %s." % self.url
         raise DownloaderError(self.url, self.local, msg)

      try:
         return self._openzipfile(zipfn, memberfn)
      except Exception as e:
         msg = "Error extracting %s from %s: %s" % (memberfn, zipfn, e)
         raise DownloaderError(self.url, self.local, msg)

   def _getfromzip(self):
      try:
         zipfn, memberfn = self.url[4:].split('?', 1)
      except Exception:
         msg = "Invalid Location: %s." % self.url
         raise DownloaderError(self.url, self.local, msg)

      src = None
      dst = None
      try:
         src = self._openzipfile(zipfn, memberfn)
         if self.fileobj:
            shutil.copyfileobj(src, self.fileobj)
            self.fileobj.seek(0)
            return self.fileobj
         else:
            dst = open(self.local, 'wb')
            shutil.copyfileobj(src, dst)
      except Exception as e:
         msg = "Error extracting %s from %s: %s" % (memberfn, zipfn, e)
         raise DownloaderError(self.url, self.local, msg)
      finally:
         if src:
            src.close()
         if dst:
            dst.close()
      return self.local

   def _getsslcontext(self):
      return self.options.get('sslcontext', None)

   def Get(self):
      """Retrieves file, setting the object's 'local' attribute to the location
         of the file on the local file system. May also change the value of
         the object's 'url' attribute, if the HTTP server returns a redirect.
            Returns: absolute path to the local file. Note that the return
                     value may be different than the 'local' parameter passed
                     to the constructor. When the 'url' parameter is a local
                     path or file:// URL, the location of the file is returned
                     without making a copy. This conserves space on the file
                     system.
            Raises:
               * DownloaderError on failure.
      """
      if not self.local:
         raise DownloaderError(self.url, self.local,
                               'Invalid download destination.')
      destdir = os.path.dirname(self.local)
      if not os.path.isdir(destdir):
         try:
            os.makedirs(destdir)
         except OSError as e:
            msg = "Cannot create directory %s: %s" % (destdir, e)
            raise Errors.FileIOError(destdir, msg)

      if self.url.startswith('zip:'):
         return self._getfromzip()
      else:
         return self._getfromurl()

   def Open(self):
      """Opens the url as a file object.  May change the value of the object's
         'url' attribute, if the HTTP server returns a redirect.
            Returns: a file object from which the callee can read directly
                     from the url to implement a download.
            Raises:
               * DownloaderError on failure.
      """
      if self.url.startswith('zip:'):
         return self._openfromzip()
      else:
         return self._openfromurl()

   @staticmethod
   def setEsxupdateFirewallRule(enabled):
      """Enables or disables the esxupdate firewall rule if the url port
         matches a port that is covered by the esxupdate firewall ruleset.
      """

      if not esxcliAvailable:
         '''If the esxclipy python binding is not available, we don't run
            on ESXi, so we don't need to try to specify any firewall rules.
         '''
         return

      # Check if firewall is unloaded. If unloaded, the default action is to
      # allow all traffic. So, no firewall ruleset change is necessary. In case
      # of failure to get the status of the firewall, we will continue to
      # attempt setting the esxupdate firewall ruleset.
      try:
         from esxutils import runCli

         cmd = ['network', 'firewall', 'get']
         result = runCli(cmd, evalOutput=True)
         if result['Loaded'] == False:
            log.info("Firewall is unloaded. Skipping esxupdate firewall "
                     "ruleset change.")
            return
      except Exception as e:
         log.warning("Unable to get firewall status. Proceeding with firewall "
                     "ruleset change. Error: %s", str(e))

      def executeEsxupdateFirewallCmd(esxcli, cmd):
         try:
            msg = 'Failed to execute "localcli %s"' % cmd
            status, output = esxcli.Execute(cmd.split(' '))
            if status != 0:
               msg = "%s with status %d, output: %s" % (
                       msg, status, ' '.join(output.splitlines()))
               log.error(msg)
               raise Exception(msg)
            else:
               log.info('Successfully executed cmd: localcli %s ', cmd)
            return
         except Exception as e:
            msg = "%s with error (%s)" % (msg, str(e))
            raise DownloaderError("Esxupdate", "Firewall Rule setting:", msg)
      try:
         cmd = '--plugin-dir /usr/lib/vmware/esxcli/int/ networkinternal ' \
               'firewall firewallRuleset set --ruleset-id esxupdate ' \
               '--service-name esxupdate --enabled %s' % enabled
         esxcli = esxclipy.EsxcliPy(True)
         executeEsxupdateFirewallCmd(esxcli, cmd)
      except DownloaderError:
         """During host upgrade, patch-the-patcher uses ESXi code from target
         image, while rest of the system still continues with older version,
         which doesn't understand the new command. Hence, there is a need to
         try with both new and old command to configure firewall ruleset. This
         behaviour is covered in try/except block below.
         """
         log.warning('Failed to execute system-owned firewall ruleset command,'
                     ' retrying with the public firewall ruleset command.')
         if enabled == 'true':
            action = 'add'
         if enabled == 'false':
            action = 'remove'
         cmd = 'network firewall ruleset client %s -r esxupdate' % action
         esxcli = esxclipy.EsxcliPy(False)
         executeEsxupdateFirewallCmd(esxcli, cmd)

   @classmethod
   def IsLinkAccesible(cls, url):
      """Returns whether a link can be accesssed.
      """
      try:
         d = cls(url)
         d.Open().close()
         return True
      except Exception:
         return False


def SetSslContext(sslcontext=None):
   """Set the SSL context that the Downloader will use.
      This accepts an ssl.SSLContext object.
      If this is set to None, an unverified context will be used.
   """
   options['sslcontext'] = sslcontext


def SetProxy(urlport=None):
   """
   Sets HTTP/FTP/HTTPS proxy, of 'url:port'
   format, such as 'http://proxy.abc.com:3128'
   If urlport is None or no argument given, the proxies option is
   deleted from options.
   """
   if urlport:
     options['proxies'] = {'http': urlport,
                           'https': urlport,
                           'ftp': urlport}
   elif 'proxies' in options:
     del options['proxies']


def SetProxies(proxies=None):
   """ Set 'http', 'https', 'ftp', 'no' proxies if provided is not None;
       otherwise, remove 'proxies' from options.
   """
   if 'proxies' in options:
      del options['proxies']

   if not proxies:
      return

   for protocol in ('http', 'https', 'ftp', 'no'):
      proxy = proxies.get(protocol + '_proxy', None)
      if proxy:
         # The module "requests" uses no_proxy than no, keep both.
         if protocol == 'no':
            options.setdefault('proxies', {})['no_proxy'] = proxy
         options.setdefault('proxies', {})[protocol] = proxy


def SetTimeout(timeout=None):
   '''
   Sets downloader option timeout value. 0/None for default behavior.
   '''
   if timeout and int(timeout) > 0:
      options['timeout'] = int(timeout)
   else:
      options.pop('timeout', 0)

def SetRetry(retry=None):
   '''Set number of times to attempt retry. Set to 0/None to only try once.
   '''
   if retry and int(retry) > 0:
      options['retry'] = int(retry)
   else:
      options['retry'] = 1

def SetRateLimit(limit=None):
   """Set the maximum rate, in bytes per second, that the downloader should
      read from a remote sources. The value may either be 0/None to indicate no
      limit, or a positive integer.
   """
   options['throttle'] = int(limit) if limit and int(limit) > 0 else 0

