#!/usr/bin/python

########################################################################
# Copyright (C) 2010-2024 VMware, Inc. All Rights Reserved             #
# VMware Confidential                                                  #
########################################################################

import logging
import os
import tempfile
import shutil
import sys
import zipfile

if sys.version_info[0] >= 3:
   from urllib.request import url2pathname
   from urllib.parse import urlparse
else:
   from urllib import url2pathname
   from urlparse import urlparse

from . import Downloader
from . import Errors
from . import Depot
from . import DepotCollection
from . import Vib

from .Utils import Misc, PathUtils

"Read, write, and extract esximage offline bundle zip files."

log = logging.getLogger('OfflineBundle')

class OfflineBundle(object):
   """Class representing an esximage offline bundle zip, with methods to scan,
      extract, and write an offline bundle zip to a file."""
   DEPOT_INDEX = 'index.xml'

   def __init__(self, bundleurl):
      """Create a new OfflineBundle instance.
         Parameters:
            * bundleurl - Either a path to an offline bundle or the full remote
                          or local URL of the depot index.xml file. Bundle file
                          name must end with '.zip'.
      """
      self._bundleurl = bundleurl
      self._dc = DepotCollection.DepotCollection()

   def Load(self, validate=False, schemaVersionCheck=False):
      ''' Read Depot metadata nodes. This is actually handled by
          DepotCollection.ConnectDepots method, but exception will be raised.
          Parameters:
            validate - If True, it enforces metadata schema validation upon
                       loading bundle.
          Exceptions:
            BundleIOError - error reading from offline bundle or a depot
      '''
      try:
         self._dc.ConnectDepots([self._bundleurl],
                                ignoreerror=False,
                                validate=validate,
                                schemaVersionCheck=schemaVersionCheck)
      except Downloader.DownloaderError as e:
         msg = 'Error in downloading files: %s' % (e)
         raise Errors.BundleIOError(self._bundleurl, msg)

   @property
   def channels(self):
      return self._dc.channels

   @property
   def vibs(self):
      return self._dc.vibs

   @property
   def profiles(self):
      return self._dc.profiles

   @property
   def vibscandata(self):
      return self._dc.vibscandata

   @property
   def solutions(self):
      return self._dc.solutions

   @property
   def manifests(self):
      return self._dc.manifests

   @property
   def baseimages(self):
      return self._dc.baseimages

   @property
   def addons(self):
      return self._dc.addons

   @property
   def bulletins(self):
      return self._dc.bulletins

   @property
   def configSchemas(self):
      return self._dc.configSchemas

   @property
   def vibExports(self):
      return self._dc.vibExports

   def ScanVibs(self):
      self._dc.ScanVibs()

   def GetBaseImage(self, releaseID):
      ''' Retrieve base image from offline bundle by provided release ID.'''
      return self._dc.GetBaseImage(releaseID)

   def GetAddon(self, releaseID):
      ''' Retrieve addon from offline bundle by provided release ID.'''
      return self._dc.GetAddon(releaseID)

   def WriteBundleZip(self, dest, checkacceptance=True,
                      partialDepotForProduct=None):
      '''Write bundle zip.
         Parameters:
            * dest            - A file path to write to.
            * checkacceptance - If True (the default), the acceptance level of
                                VIBs are validated as they are added to the
                                bundle zip.
            * partialDepotForProduct - SoftwarePlatform productLineID for which
                                       to create a partial depot.
         Exceptions:
            * BundleIOError      - Error in writing bundle zip file.
            * BundleFormatError  - If a depot metadata node or VIB is not under
                                   depot root directory.
            * VibSignatureError  - If acceptancecheck is true and acceptance
                                   level signature validation fails.
            * VibValidationError - If acceptancecheck is true and acceptance
                                   level XML schema validation fails.
      '''
      assert len(self._dc.depots) == 1, 'Only one depot is allowed'
      depotnode = self._dc.depots[0]

      try:
         bundle = zipfile.ZipFile(dest, 'w', zipfile.ZIP_DEFLATED)
      except EnvironmentError as e:
         msg = 'Error in opening file: %s' % (e)
         raise Errors.BundleIOError(dest, msg)

      depotroot = PathUtils.UrlDirname(depotnode.absurl)
      try:
         depotindex = depotnode.ToString()
         bundle.writestr(OfflineBundle.DEPOT_INDEX, depotindex)

         notificationfile = (os.path.dirname(self._bundleurl) +
                             "/notifications.zip")
         if os.path.exists(notificationfile):
            bundle.write(notificationfile, "notifications.zip")

         for vendornode in depotnode.children:
            self._AddNodeToBundle(bundle, depotroot, vendornode)
            for metanode in vendornode.children:
               self._AddNodeToBundle(bundle, depotroot, metanode, download=True)
         if partialDepotForProduct:
            vibs = self._dc.vibs.GetVibsForSoftwarePlatform(
                                                      partialDepotForProduct)
         else:
            vibs = self._dc.vibs
         for vib in vibs.values():
            self._AddVibToBundle(bundle, depotroot, vib, checkacceptance)
         bundle.close()
      except EnvironmentError as e:
         bundle.close()
         os.unlink(dest)
         msg = 'Error in writing bundle %s: %s' % (dest, e)
         raise Errors.BundleIOError(dest, msg)
      except Exception:
         bundle.close()
         os.unlink(dest)
         raise

   @staticmethod
   def _AddNodeToBundle(bundle, depotroot, node, download=False):
      log.debug('Adding DepotNode [%s] from %s' % (node.META_NODE_TAG,
         node.absurl))
      if node.absurl.startswith(depotroot):
         if download:
            with tempfile.NamedTemporaryFile() as f:
               try:
                  d = Downloader.Downloader(node.absurl, local=f.name, fileobj=f)
                  localfile = d.Get()
                  bundle.write(localfile, node.absurl[len(depotroot):])
               except Downloader.DownloaderError as e:
                  log.info('Unable to download from %s: %s', node.absurl, str(e))
         else:
            bundle.writestr(node.absurl[len(depotroot):], node.ToString())
      else:
         msg = ("Node '%s' doesn't share the same root with the depot %s" % (
            node.absurl, depotroot))
         raise Errors.BundleFormatError(bundle.filename, msg)

   @staticmethod
   def _AddVibToBundle(bundle, depotroot, vib, checkacceptance=True):
      log.debug('Adding VIB %s to bundle', vib.id)
      vurl = None
      for url in vib.remotelocations:
         if url.startswith(depotroot):
            vurl = url
            break

      if vurl is None:
         msg = 'Unable to locate %s under depot %s' % (vib.id, depotroot)
         raise Errors.BundleFormatError(bundle.filename, msg)

      scheme, _, path = urlparse(vurl)[:3]
      downloaded = False
      localfile = None
      f = None
      if scheme == 'file':
         localfile = url2pathname(path)
      else:
         f = tempfile.NamedTemporaryFile()
         try:
            d = Downloader.Downloader(vurl, local=f.name, fileobj=f)
            localfile = d.Get()
            downloaded = True
         except Downloader.DownloaderError as e:
            log.info('Unable to download from %s: %s', vurl, str(e))

      if localfile is None:
         if f:
            f.close()
         msg = "Unable to get VIB %s from URL %s" % (vib.id, vurl)
         raise Errors.VibDownloadError(vurl, '', msg)

      vibobj = None
      try:
         vibobj = Vib.ArFileVib.FromFile(localfile)
         if checkacceptance:
            vibobj.VerifyAcceptanceLevel()
         vibobj.CheckPayloadDigests()

         try:
            bundle.write(localfile, vurl[len(depotroot):])
         except EnvironmentError as e:
            msg = 'Error adding VIB %s to bundle: %s' % (vib.name, e)
            raise Errors.BundleIOError(bundle.filename, msg)
      finally:
         if vibobj:
            vibobj.Close()
         if f:
            f.close()
         if downloaded and localfile is not None:
            OfflineBundle._ForceRemoveFile(localfile)

   @staticmethod
   def _ForceRemoveFile(fn):
      if os.path.isfile(fn):
         try:
            os.unlink(fn)
         except EnvironmentError as e:
            log.info('Unable to clean up temp file %s: %s' % (fn, e))


def WriteOfflineBundle(depotFilename, vendorName, vendorCode, baseimages,
                       addons, manifests, solutions, profiles, components,
                       vibs, configSchemas=None, versions=None,
                       checkAcceptance=False,
                       metaDataZipFilename='metadata.zip',
                       legacyBulletins=None, products=None,
                       partialDepotForProduct=None, vibExports=None):
   """Writes the contents into an offline bundle.

      Params:
         * depotFilename - The depot filename to write
         * vendorName - Vendor Name to write to the depot
         * vendorCode - Vendor Code to write to the depot
         * baseimages - Baseimages to write to the depot
         * addons - Addons to write to the depot
         * manifests - Hardware support manifests to write to the depot
         * solutions - Soutions to write to the depot
         * profiles - Profiles to write to the depot
         * components - Components to write to the depot
         * vibs - Vibs to write to the depot
         * configSchemas - Config schemas to write to the depot.
         * versions - The list of metadata versions for this depot
         * checkAcceptance - Check VIB acceptance levels
         * metadataZipFilename - The metadata zip filename
         * legacyBulletins - Legacy Bulletins to write to the depot
         * products - A list of strings representing supported productIds
         * partialDepotForProduct - SoftwarePlatform productLineID for which
                                    to create a partial depot
         * vibExports - VIB exports to write to the depot.
   """
   BASE_VIB = 'esx-base'
   BASE_ESXIO_VIB = 'esxio-base'
   # default set of supported platforms/productIds
   PLATFORMS = [Vib.SoftwarePlatform.PRODUCT_EMBEDDEDESX]
   VENDOR_BASE = 'vendor-index.xml'

   if not products:
      products = PLATFORMS
   else:
      products = Misc.toDepotProductList(products)

   # create a temp directory to write out metadata.zip, XML files, and
   # VIB packages
   depotDir = tempfile.mkdtemp()
   try:
      # Download all the VIBs
      for vib in vibs.values():
         localFile = os.path.join(depotDir, vib.GetRelativePath())
         try:
            Depot.VibDownloader(localFile, vib)
         except EnvironmentError as e:
            raise Errors.VibDownloadError(', '.join(vib.remotelocations),
                                          localFile, e)

      meta = Depot.MetadataNode(url=metaDataZipFilename)
      meta.vibs.FromDirectory(depotDir, ignoreinvalidfiles=True)
      if profiles:
         meta.profiles += profiles

      if components:
         for component in components.values():
            for bullId in component:
               meta.bulletins.AddBulletin(component[bullId])

      if solutions:
         meta.solutions = solutions

      if addons:
         meta.addons = addons

      if manifests:
         meta.manifests = manifests

      if baseimages:
         meta.baseimages = baseimages

      if configSchemas:
         meta.configSchemas = configSchemas

      if vibExports:
         meta.vibExports = vibExports

      platformVersions = set()
      if versions:
         platformVersions.update(versions)

      platformVersions.update([vib.version.version.versionstring
            for vib in vibs.values() if vib.name in (BASE_VIB, BASE_ESXIO_VIB)])

      for version in platformVersions:
         meta.AddPlatform(products, version, channels=[])

      if legacyBulletins:
         # XXX There is a possibility that we could overwrite a component below
         # I'm going to skip over any tesing of this because:
         # 1) We don't have good equality operators for comparisons between
         #   components and bulletins
         # 2) Legacy components only occur in patch builds and these will
         #   everntually go away
         meta.bulletins += legacyBulletins

      # TODO: Currently when we remove a component or vib
      # if we don't remove it from an image profile, we
      # will get an error when we write the image profile back to disk
      # I'm not quiet sure what the policy on this should be or if
      # the EPK should even be storing image profiles in the manner
      # we do. I'm temporarily turning off the logger for warnings
      # while we write back the image profiles to suppress this
      # problem until a solution is figured out.
      logger = logging.getLogger()
      curLevel = logger.getEffectiveLevel()
      logger.setLevel(logging.ERROR)
      meta.WriteMetadataZip(os.path.join(depotDir, metaDataZipFilename))
      logger.setLevel(curLevel)

      # Create vendor-index.xml
      vendorIndex = Depot.VendorIndex(name=vendorName,
                                      code=vendorCode,
                                      indexfile=VENDOR_BASE,
                                      children=[meta])

      path = os.path.join(depotDir, VENDOR_BASE)
      try:
         with open(path, 'wb') as vendorFile:
            vendorFile.write(vendorIndex.ToString())
      except IOError as e:
         raise Errors.BundleIOError(path,
                                    'Error writing out vendor-index.xml: %s'
                                    % e)

      # create index.xml
      depotIndex = Depot.DepotIndex(children=[vendorIndex])
      path = os.path.join(depotDir, 'index.xml')
      try:
         with open(path, 'wb') as indexFile:
            indexFile.write(depotIndex.ToString())
      except IOError as e:
         raise Errors.BundleIOError(path,
                                    'Error writing out index.xml: %s' % e)

      offlineBundle = OfflineBundle(depotDir)
      offlineBundle.Load()
      offlineBundle.WriteBundleZip(depotFilename,
                              checkacceptance=checkAcceptance,
                              partialDepotForProduct=partialDepotForProduct)

   finally:
      shutil.rmtree(depotDir)


def CreatePartialOfflineDepot(profile, platform, vibdownloadfn=Depot.VibDownloader):
   """Creates partial depot for the platform and returns the created depot.
      Returns None if the profile is meant for only 1 platform.

      The caller has to cleanup the depot created.

      Parameters:
         * profile - An instance of ImageProfile, with the vibs attribute
                     containing a VibCollection with all of the VIBs in
                     vibIDs with the sourceurl attribute populated.
         * platform - SoftwarePlatform productLineID of the partial depot to
                      be created.
         * vibdownloadfn - Function for downloading a vib object. The function
                           signature should be
                           fn(destfilepath, vibobj, extraArgs=None)
   """
   if len(profile.GetSoftwarePlatforms(
                             fillDefaultValue=False, baseEsxOnly=True)) > 1:
      try:
         esxioDepot = tempfile.NamedTemporaryFile(delete=False)
         with tempfile.TemporaryDirectory() as tmpDir:
            Depot.DepotFromImageProfile(profile, tmpDir,
               vendor='VMware, Inc.',
               vendorcode='vmw',
               generateRollupBulletin=False,
               vibdownloadfn=vibdownloadfn)

            offlineBundle = OfflineBundle(tmpDir)
            offlineBundle.Load()
            offlineBundle.WriteBundleZip(esxioDepot.name,
               checkacceptance=False, partialDepotForProduct=platform)

      except Exception:
         if esxioDepot and os.path.isfile(esxioDepot.name) \
                       and not esxioDepot.closed:
            esxioDepot.close()
            os.unlink(esxioDepot.name)
         raise
      else:
         esxioDepot.close()
         return esxioDepot.name
   return None


if __name__ == '__main__':
   logging.basicConfig(level=logging.DEBUG)

   metaurl = sys.argv[1]
   dest = sys.argv[2]
   ob = OfflineBundle(metaurl)
   ob.Load()
   ob.WriteBundleZip(dest)
