########################################################################
# Copyright (C) 2008-2021 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:
   import urllib.request
   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.
                      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)
      import ssl

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

      # XXX: create https handler with no verification
      https_handler = HTTPSHandler(context=ssl._create_unverified_context())
      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.warning('File not found (Error 404). Download will not '
                              'be retried')
                  raise
            _checkRetry(retry, retries, err)
         except Exception as e:
            _checkRetry(retry, retries, e)

   def _generate_url_request(self):
      ''' Generate a get request for the URL.
      '''
      import requests
      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
          requests.utils.should_bypass_proxies(self.url, proxies['no_proxy'])):
         proxies = None

      if certExists and proxies:
         req = requests.get(self.url, stream=True, allow_redirects=True,
                            proxies=proxies, verify=CA_CERT_DIR)
      elif proxies:
         req = requests.get(self.url, stream=True, allow_redirects=True,
                            proxies=proxies)
      elif certExists:
         req = requests.get(self.url, stream=True, allow_redirects=True,
                            verify=CA_CERT_DIR)
      else:
         req = requests.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 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
      if enabled == 'true':
         action = 'add'
      if enabled == 'false':
         action = 'remove'
      cmd = 'network firewall ruleset client %s -r esxupdate' % action
      msg = 'Failed to execute "esxcli %s"' % cmd
      try:
         esxcli = esxclipy.EsxcliPy(False)
         status, output = esxcli.Execute(cmd.split(' '))
         if status != 0:
            msg = "%s with status %s" % (msg, status)
            log.info(msg)
         return
      except Exception as e:
         msg = "%s with error (%s)" % (msg, str(e))
         raise DownloaderError("Esxupdate", "Firewall Rule setting:", msg)


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

