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

"""This module defines classes for reading and writing a depot's hierarchy of
   metadata files, including index.xml and vendor-index.xml.
"""

import logging
import os
import shutil

from . import Vib
from . import Errors
from . import Metadata
from . import VibCollection
from .Utils import ArFile

try:
   from . import Downloader
   HAVE_DOWNLOADER = True
except ImportError:
   HAVE_DOWNLOADER = False

from . import Bulletin
from .Database import TarDatabase
from .Utils import PathUtils, XmlUtils

# Use the XmlUtils module to find ElementTree.
etree = XmlUtils.FindElementTree()

log = logging.getLogger(__name__)

ESX_DEPOT_CONTENT_NAME = "VMware ESX"
ESX_DEPOT_CONTENT_TYPE = "http://www.vmware.com/depotmanagement/esx"

class DepotTreeNode(object):
   """Represents one node in a tree of depot metadata files.
      Each node represents one XML metadata file which points at children
      metadata files.
      This class is meant to be an abstract base class and inherited by
      each level of the tree hierarchy.

      Class Attributes:
         * baseurl     - Absolute URL of parent node, or None
         * url         - Absolute or relative url, passed into constructor
         * absurl      - Absolute URL of this node
         * children    - A list of child nodes, each of them a DepotTreeNode.
   """
   # The XML tag name for the entire metadata file.
   META_NODE_TAG = "metadata"

   def __init__(self, url=None, baseurl=None, children=None):
      """Constructs a new DepotTreeNode instance.  Computes the absolute remote
         URL and local path based on a relative url, if necessary.
         Parameters:
            * url       - Either an absolute URL pointing at this node, or
                          a relative path of the form "dir1/dir2".  If this is
                          the top of the tree, an absolute URL should be used.
            * baseurl   - The absolute URL of the parent node.  Used to compute
                          the absolute URL of this node if 'url' is relative.
            * children  - A list of child nodes, which should be some subclass
                          of DepotTreeNode.
         Returns:
            A new instance of DepotTreeNode.
      """
      self.baseurl = baseurl
      self.url = url
      self.absurl = None
      self.children = list()
      #
      # Use AddChild to construct the children property so that any duplicates
      # can be merged. This is used to merge multiple <metadata> nodes that
      # refer to the same metadata.zip, for example.
      #
      if children is not None:
         for child in children:
            self.AddChild(child)
      if url:
         self.absurl = PathUtils.UrlJoin(baseurl, url)

   @classmethod
   def _XmlToKwargs(cls, xml, url=None):
      raise NotImplementedError("This method must be instantiated by a subclass.")

   @classmethod
   def FromXml(cls, xml, url=None):
      """Constructs a new DepotTreeNode instance from XML representing the
         entire metadata file.  This method will not be implemented for the
         leaf nodes in a tree, for which there are no individual metadata files.
         Parameters:
            * xml      - Either a string or an ElementTree instance containing
                         the depot node XML
            * url      - The absolute URL of this XML file
         Returns:
            A new DepotTreeNode instance
         Raises:
            MetadataFormatError - badly formatted or unparseable metadata XML.
      """
      if not etree.iselement(xml):
         try:
            xml = etree.fromstring(xml)
         except Exception as e:
            raise Errors.MetadataFormatError(None,
               "Could not parse %s XML data: %s." % (cls.__class__.__name__, str(e)))

      kwargs = cls._XmlToKwargs(xml, url=url)
      return cls(**kwargs)

   @classmethod
   def FromChildXmlNode(cls, xml, baseurl=None):
      """Constructs a new DepotTreeNode instance from XML representing a child
         node inside another metadata file.  This method will not be implemented
         by the top node of a tree.
         Parameters:
            * xml  - An ElementTree instance representing a child node element
            * baseurl - The absolute URL representing the location of parent
         Returns:
            A new DepotTreeNode instance
         Raises:
            MetadataFormatError - child node is missing required elements, or
                                  other parsing error
      """
      raise NotImplementedError("This method must be instantiated by a subclass.")

   def ToXml(self):
      """Serializes this DepotTreeNode instance out to XML for an entire
         metadata file.
         Returns:   An ElementTree instance.
         Exceptions:
            None
      """
      metanode = etree.Element(self.META_NODE_TAG)
      for child in self.children:
         nodes = child.ToChildXmlNode()
         for node in nodes:
            metanode.append(node)
      return metanode

   def ToString(self):
      """Writes out this DepotTreeNode instance as a string.
         Returns:   None
      """
      try:
         # pretty_print is supported by lxml only
         return etree.tostring(self.ToXml(), pretty_print=True)
      except Exception:
         return etree.tostring(self.ToXml())

   def ToChildXmlNode(self):
      """Serializes this DepotTreeNode into child node(s) for inclusion
         in a metadata file.
         Returns:   A list of ElementTree instances.
         Exceptions:
            None
      """
      raise NotImplementedError("This method must be instantiated by a subclass.")

   def GetChildFromUrl(self, absurl):
      """Finds a child node with a matching URL.
         Parameters:
            * absurl - The absolute URL for the child node to find
         Returns:
            An instance of DepotTreeNode, or None if no matching child is
            found.
      """
      for child in self.children:
         if child.absurl == absurl:
            return child
      return None

   def AddChild(self, childnode):
      """Adds a child node to this node.  If the child node already exists
         (an existing child node has matching absurl), then the nodes are
         merged, with preference given to attributes in childnode.  An example
         of where the merging is necessary is that an instance of DepotTreeNode
         may be created with FromChildNodeXml() from the parent node; subsequently
         the absolute URL of the child metadata is determined, and that metadata
         file is then parsed using FromFile().  The two instances need to be
         merged.

         The childnode is not modified.

         Parameters:
            * childnode - An instance of DepotTreeNode or subclass.
         Returns:
            The node updated or added.
      """
      mychild = self.GetChildFromUrl(childnode.absurl)
      if mychild:
         newchild = mychild + childnode
         # Make sure we replace the original reference in children
         self.children[self.children.index(mychild)] = newchild
         return newchild
      else:
         self.children.append(childnode)
         return childnode

   ATTRS_TO_COPY = ("baseurl", "url")
   def __add__(self, other):
      """Creates a new instance based on the merging of this instance
         and other.  Attributes from other are given a preference.
      """
      kwargs = {'children': self.children}
      for attr in self.ATTRS_TO_COPY:
         kwargs[attr] = getattr(self, attr)
      #
      # It is quite likely that parental information is only present
      # in instances created using FromChildNodeXml.
      # Also, instead of doing a deepcopy or shallow copy, copying
      # attributes allows us to preserve attributes in self which
      # are not in other.
      for attr in self.ATTRS_TO_COPY:
         otherval = getattr(other, attr)
         if otherval:
            kwargs[attr] = otherval
      #
      # Merge the children. This recurses down the hierarchy
      # if needed.
      kwargs['children'].extend(other.children)

      return self.__class__(**kwargs)


class DepotIndex(DepotTreeNode):
   """This class represents an index.xml file at the top of a depot tree.
      It points to one or more vendor-index.xml files.
   """
   META_NODE_TAG = "vendorList"

   @classmethod
   def _XmlToKwargs(cls, xml, url=None):
      kwargs = {'children': []}
      for vendornode in xml.findall('vendor'):
         child = VendorIndex.FromChildXmlNode(vendornode, baseurl=url)
         kwargs['children'].append(child)
      kwargs['url'] = url
      return kwargs


class VendorIndex(DepotTreeNode):
   """This class represents a vendor-index.xml file, and it also represents
      the data in each <vendor> node of index.xml.
      A vendor-index.xml file contains a list of metadata.zip locations and
      properties.  Each metadata.zip file belongs to one or more channels.

      Interesting attributes:
         * channels - A dict whose keys are the channel names from the individual
                      <metadata> nodes.  The value for each channel key is an
                      instance of DepotChannel.  If a metadata node has no
                      channel listed, it goes under the 'default' key.
   """
   META_NODE_TAG = "metaList"
   ATTRS_TO_COPY = ("baseurl", "url", "name", "code", "indexfile",
                    "patchUrl", "vibUrl",
                    "relativePath", "contentName", "contentType")

   def __init__(self, **kwargs):
      """Initialize a VendorIndex object.
         Unique Parameters:
            * name      - The name of the vendor, ex "VMware, Inc."
            * code      - Shortened vendor code, ex "VMW"
            * indexfile - name of the vendor-index.xml file
            * patchUrl  - relative or absolute URL of directory containing indexfile (deprecated)
            * vibUrl    - Optional URL to root of VIBs if different than patchUrl
            * relativePath   - relative path of directory containing indexfile
            * contentName    - readable name of the depot
            * contentType    - internal id of the depot schema/types/structure
      """
      self.name      = kwargs.pop('name', '')
      self.code      = kwargs.pop('code', '')
      self.indexfile = kwargs.pop('indexfile', '')
      self.patchUrl  = kwargs.pop('patchUrl', '')
      self.vibUrl    = kwargs.pop('vibUrl', None)
      self.relativePath   = kwargs.pop('relativePath', '')
      self.contentName    = kwargs.pop('contentName', ESX_DEPOT_CONTENT_NAME)
      self.contentType    = kwargs.pop('contentType', ESX_DEPOT_CONTENT_TYPE)
      if 'url' not in kwargs or not kwargs['url']:
         if self.relativePath:
            kwargs['url'] = '/'.join([self.relativePath.rstrip('/'), self.indexfile])
         else:
            kwargs['url'] = self.indexfile
      DepotTreeNode.__init__(self, **kwargs)

   @classmethod
   def FromChildXmlNode(cls, xml, baseurl=None):
      kwargs = {}
      for attr in ("name", "code", "indexfile", "relativePath"):
         node = xml.find(attr)
         if node is None:
            raise Errors.MetadataFormatError(None,
               "Element %s was expected, but not found" % (attr))
         kwargs[attr] = node.text and node.text.strip() or ''

      for attr in ("vibUrl", "patchUrl"):
         val = xml.findtext(attr, None)
         if val:
            kwargs[attr] = val.strip()

      content = xml.find("content")
      if content is not None:
         val = content.findtext("name", None)
         if val is None:
            val=""
         kwargs["contentName"] = val.strip()

         val = content.findtext("type", None)
         if val:
            kwargs["contentType"] = val.strip()

      kwargs['baseurl'] = baseurl

      return cls(**kwargs)

   def ToChildXmlNode(self):
      node = etree.Element("vendor")
      for attr in ("name", "code", "indexfile", "patchUrl"):
         etree.SubElement(node, attr).text = getattr(self, attr)
      if self.vibUrl:
         etree.SubElement(node, "vibUrl").text = self.vibUrl
      if self.relativePath is not None:
         etree.SubElement(node, "relativePath").text = self.relativePath
      if (not self.contentName is None) or (not self.contentType is None):
         content=etree.SubElement(node, "content")
         if not self.contentName is None:
            etree.SubElement(content, "name").text=self.contentName
         if not self.contentType is None:
            etree.SubElement(content, "type").text=self.contentType

      return [node]

   @classmethod
   def _XmlToKwargs(cls, xml, url=None):
      kwargs = {'children': []}
      for metanode in xml.findall('metadata'):
         child = MetadataNode.FromChildXmlNode(metanode, baseurl=url)
         kwargs['children'].append(child)

      text = xml.findtext('vibUrl', None)
      if text:
         kwargs['vibUrl'] = text.strip()
      text = xml.findtext('relativePath', None)
      if text:
         kwargs['relativePath'] = text.strip()
      content = xml.find('content')
      if content:
         val = content.findtext('name',None)
         if not val is None:
            kwargs['contentName'] = val.strip()
         val = content.findtext('type',None)
         if not val is None:
            kwargs['contentType'] = val.strip()

      kwargs['url'] = url
      return kwargs

   def ToXml(self):
      xml = DepotTreeNode.ToXml(self)
      if self.vibUrl:
         etree.SubElement(xml, 'vibUrl').text = self.vibUrl
      return xml

   def GetChannels(self):
      """Returns the metadata for each available channel.
         Parameters:
            None
         Returns:
            A dict whose keys are the channel name strings from each
            metadata node.  The value for each key is an instance of
            DepotChannel.  If a MetadataNode does not have any channels
            listed, it goes under a channel by the name 'default'.
      """
      channels = {}
      for meta in self.children:
         metamap = meta.GetChannelPlatformMap()
         for channelname in metamap:
            if channelname not in channels:
               channel = DepotChannel(channelname, self.absurl,
                                      metadatas = [meta],
                                      vendorindex = self)
               channels[channelname] = channel
            else:
               channels[channelname].metadatas.append(meta)
      return channels

   #
   # There aren't that many metadatas per vendor file, so this should be
   # adequate as a property.  If performance really becomes a concern,
   # we can cache the results and compute them at the end of __init__
   # and after AddChild().
   #
   channels = property(GetChannels)


class DepotChannel(object):
   """Represents one channel in a depot.
      Attributes:
         * channelId - the globally unique channel ID
         * name      - name of channel, only unique within one vendor-index.xml
         * vendorIndexUrl - the URl of the vendor-index file containing this
         * vendorindex - the VendorIndex instance containing this channel
         * metadatas - a list of MetadataNode instances
   """
   def __init__(self, name, vendorIndexUrl, metadatas=None, vendorindex=None):
      self.name = name
      self.vendorIndexUrl = vendorIndexUrl
      self.vendorindex = vendorindex
      if metadatas is None:
         metadatas = list()
      self.metadatas = metadatas

   def _getUniqueId(self):
      return str((self.vendorIndexUrl, self.name))

   channelId = property(_getUniqueId)

   def __hash__(self):
      return hash(self.channelId)

   def __eq__(self, other):
      return self.channelId == other.channelId

   def _HasReleaseUnit(self, typeName, releaseID):
      """ Check that this depot channel contains a release unit with the
          provided release unit type and release unit ID.
      """
      for meta in self.metadatas:
         if meta.HasReleaseUnit(typeName, releaseID):
            return True
      return False

   def RemoveMatchedReleaseUnits(self, releaseUnits):
      """ Remove release units existing both in the provided release unit
          set and this channel from releaseUnits.

          If one or more release units are removed, return True;
          otherwise, False.
      """
      found = set()
      for typeName, relID in releaseUnits:
         if self._HasReleaseUnit(typeName, relID):
            found.add((typeName, relID))

      releaseUnits.difference_update(found)
      return bool(found)

   def RemoveMatchedComponentIDs(self, knownCompIDs):
      """ Remove components existing both in the provided component ID
          set and this channel from knownCompIDs.

          If one or more components are removed, return True; otherwise, False.
      """
      found = False
      for meta in self.metadatas:
         cc = Bulletin.ComponentCollection(meta.bulletins,
                                           ignoreNonComponents=True)
         foundIDs = knownCompIDs.intersection(cc.GetComponentNameIds())
         if foundIDs:
            found = True
            knownCompIDs.difference_update(foundIDs)
      return found

class MetadataNode(DepotTreeNode, Metadata.Metadata):
   """A leaf node in the depot tree, representing a metadata.zip file and its
      associated (product, version, locale) groupings.
      A metadata.zip can belong to one or more "channels", identified by name.
      Attributes:
         * platforms - A list of supported platforms and channels.  Each member
                       is a tuple: (product, version, locale, channels); where
                       the first three defines a supported platform, and channels
                       is a list of channel names (or []) supported by this metadata
                       for that given platform.
   """
   def __init__(self, **kwargs):
      """Constructor for a MetadataNode.
         Unique Parameters:
            * productId - The Product ID for which the metadata contains packages
            * version   - The version string for the supported product
            * locale    - The locale for the supported product
            * channels  - A list of strings each representing a channel to which
                          this metadata belongs
         Returns:
            A new MetadataNode instance
      """
      productId = kwargs.pop('productId', '')
      version   = kwargs.pop('version', '')
      locale    = kwargs.pop('locale', '')
      channels  = kwargs.pop('channels', [])
      self.platforms = list()
      if productId:
         self.AddPlatform(productId, version, locale, channels)
      DepotTreeNode.__init__(self, **kwargs)
      Metadata.Metadata.__init__(self)

   def __add__(self, other):
      #
      # Overload the add method to support merging of the platforms list.
      # This allows different MetadataNodes which have been created from
      # vendor-index.xml's defining different platforms for the same
      # metadata.zip file to be merged together when vendor-index.xml is
      # created.
      #
      new = DepotTreeNode.__add__(self, other)
      new.platforms.extend(self.platforms)
      new.platforms.extend(other.platforms)
      return new


   def AddPlatform(self, productId, version, locale='', channels=[]):
      if not channels:
         channels = ['default']
      newtuple = (productId, version, locale, channels)
      if newtuple not in self.platforms:
         self.platforms.append(newtuple)


   def _parseExtraFile(self, filename, xmltext):
      if filename != 'vendor-index.xml':
         return "false"

      xml = etree.fromstring(xmltext)
      for meta in xml.findall('metadata'):
         p   = meta.find("productId")
         v   = meta.find("version")
         l   = meta.find("locale")
         url = meta.find("url")
         self.url = url.text
         chnls = meta.findall("channelName")
         nms = list()
         for c in chnls:
            nms.append(c.text)
         if l.text==None:
            l.text=''
         self.AddPlatform(p.text, v.text, l.text, nms)

      return "true"


   def _writeExtraMetaFiles(self, stagedir):
      metanode = etree.Element("metaList")
      nodes = self.ToChildXmlNode()
      for node in nodes:
         metanode.append(node)

      XmlUtils.IndentElementTree(metanode)

      vixml = os.path.join(stagedir, 'vendor-index.xml')
      tree = etree.ElementTree(element=metanode)
      try:
         tree.write(vixml)
      except Exception as e:
         msg = 'Failed to write vendor-index.xml file: %s' % e
         raise Errors.MetadataBuildError(msg)


   def GetChannelPlatformMap(self):
      """Returns a mapping of channel names to supported platforms.
         Returns:
            A dict of the form {<channelName>: [(p1,v1,l1), (p2, v2, l2), ...]}.
            Each key is the channel name found in <channelName>, or 'default'
            for platforms defined with no channel name.
            Each value is a list of (productID, version, locale) tuples.
      """
      chanmap = {}
      for prod, ver, locale, channels in self.platforms:
         for channelname in channels:
            chanmap.setdefault(channelname, []).append((prod, ver, locale))
      return chanmap

   @classmethod
   def FromChildXmlNode(cls, xml, baseurl=None):
      kwargs = {}
      for attr in ("productId", "version", "locale", "url"):
         node = xml.find(attr)
         if node is None:
            raise Errors.MetadataFormatError(None,
               "Element %s was expected, but not found" % (attr))
         kwargs[attr] = node.text and node.text.strip() or ''

      kwargs['channels'] = []
      for channelNode in xml.findall('channelName'):
         kwargs['channels'].append(channelNode.text.strip())

      kwargs['baseurl'] = baseurl

      return cls(**kwargs)

   def ToChildXmlNode(self):
      nodes = []
      if self.notifications:
         node = etree.Element("notification")
         etree.SubElement(node, "url").text = "notifications.zip"
         nodes.append(node)
      for product, ver, locale, channels in self.platforms:
         node = etree.Element("metadata")
         etree.SubElement(node, "productId").text = product
         etree.SubElement(node, "version").text = ver
         etree.SubElement(node, "locale").text = locale
         etree.SubElement(node, "url").text = self.url
         for channel in channels:
            etree.SubElement(node, "channelName").text = channel
         nodes.append(node)

      return nodes


DEPOT_PRODUCT = 'embeddedEsx'
DEPOT_VERSIONS = ['5.*']

def VibDownloader(destfile, vibobj, checkdigests=False, extraArgs=None):
   fn = None
   errors = []
   destdir = os.path.dirname(destfile)
   if not os.path.exists(destdir):
      os.makedirs(destdir)
   disable404Retry = False
   if extraArgs is not None and extraArgs.get('isReservedVib', None):
      log.info('For reserved vibs, download retry shall be disabled on 404 '
               'error.')
      disable404Retry = True
   for remoteurl in vibobj.remotelocations:
      try:
         d = Downloader.Downloader(remoteurl, local=destfile,
                                   disable404Retry=disable404Retry)
         fn = d.Get()

         # TODO: Optimize this in the future to do the checksums as the file
         # is being downloaded.
         if checkdigests:
            arvibobj = Vib.ArFileVib.FromFile(fn)
            arvibobj.CheckPayloadDigests()
         break
      except Downloader.DownloaderError as e:
         log.info("Skipping URL %s for VIB %s: %s",
                  remoteurl, vibobj.id, str(e))
         errors.append(str(e))
         continue
   if not fn:
      raise Errors.VibDownloadError(', '.join(vibobj.remotelocations), destfile,
                                    "Unable to download from any URLs: %s"
                                    % (', '.join(errors)))
   if os.path.normpath(destfile) != os.path.normpath(fn):
      shutil.copy2(fn, destfile)

def GenerateVib(destfile, vibobj, checkdigests=False, extraArgs=None):
   """Generates a VIB in the system at destfile location. If the VIB is
      installed, re-create it and verify the checksum of the recreated VIB.
      For a reserved VIB, copy it from the reserved VIB cache, raise an
      exception if the VIB is not present.
      Recreation of VIB is done using orig descriptor with payloads from 3
      possible sources:
         1) ISO extract dir during an VUM upgrade, for both incoming bootbank
            and locker VIBs.
         2) Current bootbank, for bootbank VIBs installed on the system.
         3) Locker partition, for locker VIB installed on the system.
   """
   # TODO: Handle error message handling to be presented on VC side.
   from .ImageManager.HostSeeding import (HOST_SEED_DIR_NAME,
      getBootBankPayloadPath, getIsoUpgradePayloadPath, getLockerPayloadPath,
      VibNotInCacheError)

   tmpWorkDir = os.path.join(HOST_SEED_DIR_NAME, 'generateVib')

   def makedirs(dirPath):
      try:
         os.makedirs(dirPath, exist_ok=True)
      except OSError as e:
         raise Errors.VibRecreateError(vibobj.id, "Cannot create directory "
                                       "%s: %s" % (dirPath, str(e)))

   def rmtree(dirPath):
      try:
         shutil.rmtree(dirPath, ignore_errors=True)
      except OSError as e:
         log.warning('Failed to remove folder %s: %s', dirPath, str(e))

   def editVibHeader(destfile):
      '''
      Fallback mechanism for vibs generated using python2.x or ar command.
      Edits the header of the vib to mimick the tools used to generate
      them.
      '''

      # TODO: This function must be removed when we have all vibs generated
      # using Python3.x.
      with open(destfile, 'rb+') as f:
         ar = ArFile.ArFile(fileobj=f)
         descInfo, _ = ar.next()
         signInfo, _ = ar.next()
         payloadInfo, _ = ar.next()

         fileInfo = [descInfo, signInfo, payloadInfo]

         isEsxUI = (payloadInfo.filename == 'esx-ui')

         for info in fileInfo:
            # esx-ui vib is built using ar (archive) command.
            # ar -D command adds a file to the archive with a file mode of 644
            # ar command also puts a trailing '/' at the end of the filename
            if isEsxUI:
               info.mode = '644'
               if not info.filename.endswith('/'):
                  info.filename += '/'
            else:
               # vibs built using Python2.x
               info.mode = '0'

            # Refer to ArFile.py for header format and indices.
            # info.offset() gives the index at which the file content starts. Walk
            # backwards until the start of the header and replace it.
            f.seek(info.offset-60)
            header = '%-16.16s%-12.12s%-6.6s%-6.6s%-8.8s%-10.10s%-2.2s' % (
                     info.filename, info.timestamp, info.uid,
                     info.gid, info.mode, info.size,
                     ArFile.AR_FILEMAGIC_STR)
            f.write(header.encode('utf-8'))


   def checkVibChecksum(raiseException=False):
      """Verify checksum of recreated vib as an early check before offline
         bundle.
         If the vib object doesn't contain checksum, then we will skip
         verification for it.

         Parameters:
            raiseException - Indicates whether an exception will be raised
                             on checksum mismatch
      """
      if not vibobj.checksum.checksum:
         log.debug("VIB %s checksum is not available, skip verification",
                   vibobj.id)
      else:
         try:
            checksum = Vib.Checksum(
                          "sha-256",
                          VibCollection._getdigest(destfile, 'sha256'))
         except Exception as e:
            raise Errors.VibChecksumError(vibobj.id, "Cannot calculate VIB "
                                          "checksum: %s" % str(e))
         if vibobj.checksum != checksum:
            if not raiseException:
               log.warning("Checksum did not match. vib may have been "
                           "created using older version of tools, retrying")
               return False
            else:
               raise Errors.VibChecksumError(vibobj.id, "VIB checksum does not "
                                          "match: calculated %s, expected: %s" %
                                          (checksum.checksum,
                                           vibobj.checksum.checksum))
      return True

   csDir = extraArgs.get('configSchemaDir', None)

   def verifyVibChecksumWithWorkaround():
      if checkVibChecksum() != True:
         editVibHeader(destfile)
         checkVibChecksum(raiseException=True)

   def stageConfigSchema(vib, vibPath, tmpDir):
      """Copy or extract config schema of a VIB into configSchemaDir.
      """
      csTag = vib.GetConfigSchemaTag()
      if not csTag or not csDir:
         # No need/nothing to stage.
         return

      csPath = os.path.join(os.sep, csTag.payloadFilePath)
      destPath = os.path.join(csDir, csTag.schemaFileName)
      if os.path.isfile(csPath):
         # Installed VIB, copy over.
         try:
            shutil.copy2(csPath, destPath)
         except shutil.Error as e:
            raise Errors.VibRecreateError(vib.id, "Failed to copy config schema"
               " from %s to %s: %s" % (csPath, destPath, e))
      else:
         # Extract from the VIB.
         vibObj = None
         try:
            # Open the VIB fresh to make sure we have proper payload access.
            vibObj = Vib.ArFileVib.FromFile(vibPath)
            cs = Vib.getConfigSchema(vibObj, None, tmpDir)
            cs.WriteFile(destPath)
         except Exception as e:
            raise Errors.VibRecreateError(vib.id, "Failed to extract config "
               "schema of VIB %s: %s" % (vib.id, e))
         finally:
            if vibObj:
               vibObj.Close()

   # If hostImage is not provided in extraArgs, do not proceed further
   if not extraArgs:
      raise Errors.VibRecreateError(vibobj.id, "Required argument not provided")

   hostImage = extraArgs.get('hostImage', None)
   if hostImage is None:
      raise Errors.VibRecreateError(vibobj.id, "HostImage object not provided")

   makedirs(tmpWorkDir)

   # Create the directory for holding the vib
   destdir = os.path.dirname(destfile)
   makedirs(destdir)

   # Indicates if we are generating a VIB from ISO extract dir in VUM upgrade.
   isoDir = extraArgs.get('isoDir', None)

   # Set of reserved VIBs.
   resVibIds = extraArgs.get('resVibIds', set())

   if not isoDir and vibobj.id in resVibIds:
      # If we are generating a reserved VIB, it should have a copy in the
      # reserved VIB cache.
      # TODO: create a separate error class for reserved VIB errors.
      if 'resVibCache' not in extraArgs:
         raise Errors.VibRecreateError(vibobj.id,
            "Reserved VIB cache object is not provided")

      resVibCache = extraArgs['resVibCache']
      if not resVibCache:
         raise Errors.VibRecreateError(vibobj.id,
            "Failed to generate reserved VIB %s: no reserved VIB cache "
            "available" % vibobj.id)

      try:
         resVibPath = resVibCache.getVibLocation(vibobj.id)
         shutil.copy2(resVibPath, destfile)
      except VibNotInCacheError:
         raise Errors.ReservedVibExtractError(vibobj.id,
            "Failed to add reserved VIB %s: not found in the reserved "
            "VIB cache storage" % vibobj.id)
      except OSError as e:
         raise Errors.ReservedVibExtractError(vibobj.id,
            "Failed to copy reserved VIB %s to %s: %s"
            % (vibobj.id, destfile, str(e)))
      except Exception as e:
         # Unexpected error
         log.exception('Failed to fetch VIB from reserved VIB cache')
         raise Errors.ReservedVibExtractError(vibobj.id,
            "Failed to add reserved VIB %s due to unhandled error: %s"
            % (vibobj.id, str(e)))

      verifyVibChecksumWithWorkaround()

      stageConfigSchema(vibobj, destfile, tmpWorkDir)
      return

   descxmlOrig = vibobj.GetOrigDescriptor()
   if not descxmlOrig:
      raise Errors.VibRecreateError(vibobj.id, "descriptor.xml not found")

   # Create a ArFileVib instance and pass through the payloads and add to
   # the payloadList which will be used later when payloads are added to vib
   arVib = Vib.ArFileVib.FromXml(descxmlOrig)
   if not arVib.payloads:
      raise Errors.VibRecreateError(vibobj.id, "No payloads present")

   payloadList = []
   for p in arVib.payloads:
      bootorder = None
      if p.bootorder:
         bootorder = p.bootorder
      payloadList.append((p.name, p.payloadtype, p.bootorder))

   # Create a new BaseVib instance which will hold all the information for
   # vib which is being recreated.
   sign = vibobj.GetSignature()
   try:
      vib = Vib.BaseVib.FromXml(descxmlOrig, validate=False, signature=sign)
   except Errors.VibValidationError as e:
      raise Errors.VibRecreateError(vibobj.id, "Error validating "
                                    "descriptor.xml: %s" % e.msg)
   except Errors.VibFormatError as e:
      raise Errors.VibRecreateError(vibobj.id, "Incorrect descriptor format: "
                                    "%s" % e.msg)
   vib.setvibtype(vibobj.vibtype)
   vib.SetOrigDescriptor(descxmlOrig)

   if isoDir:
      # VUM ISO upgrade case, image profile of the ISO is required.
      dbPath = hostImage.TryLowerUpperPath(isoDir, 'imgdb.tgz')
      if not dbPath:
         raise Errors.VibRecreateError(vibobj.id, "Failed to locate image "
            "database in ISO extract dir %s" % isoDir)
      isoDb = TarDatabase(dbPath)
      isoDb.Load()
      isoProfile = isoDb.profile

   # Pass through payloadList and add the payloads one by one to Vib.
   # Check if payload type is valid. Local file path of payload is picked
   # and checked if the payload is actually staged there. If it passes all
   # the checks, we will add the payload to the vib.
   for name, payloadType, bootorder in payloadList:
      # Payload class only warns for bad payload type, so we need to check.
      if payloadType not in Vib.Payload.PAYLOAD_TYPES:
         raise Errors.VibRecreateError(vibobj.id, "Payload type must be one of"
                                       ": %s" % str(Vib.Payload.PAYLOAD_TYPES))
      try:
         # bootorder can legally be zero, which evaluates to false.
         if bootorder is not None:
            pObj = Vib.Payload(name, payloadType, bootorder=bootorder)
         else:
            pObj = Vib.Payload(name, payloadType)

         if vibobj.name == 'esx-base':
            # esx-base's txt-mle checksum of payload 'b' needs to be copied
            # from the original VIB since it cannot be re-calculated without
            # objdump.
            for payload in vibobj.payloads:
               # Can't use IterPayload() without file object.
               for checksum in payload.checksums:
                  if checksum.verifyprocess == 'txt-mle':
                     pObj.checksums.append(checksum)
                     break

      except ValueError as e:
         rmtree(tmpWorkDir)
         raise Errors.VibRecreateError(vibobj.id, "%s" % e.msg)

      try:
         if isoDir:
            # VUM ISO upgrade directory (all VIBs).
            stagePath = getIsoUpgradePayloadPath(vibobj, name, hostImage,
                                                 isoDir, isoProfile)
         elif vibobj.vibtype == Vib.BaseVib.TYPE_BOOTBANK:
            # Bootbank VIB on disk.
            stagePath = getBootBankPayloadPath(vibobj, name, hostImage,
                                               tmpWorkDir)
         else:
            # Locker VIB on disk.
            stagePath = getLockerPayloadPath(vibobj, name, hostImage,
                                             tmpWorkDir)
      except Exception as e:
         rmtree(tmpWorkDir)
         msg = ('Cannot find or stage payload %s type %s of VIB %s: %s'
                % (name, payloadType, vibobj.id, str(e)))
         raise Errors.VibRecreateError(vibobj.id, msg)

      if stagePath is None:
         continue

      try:
         vib.AddPayload(pObj, stagePath, tmpDir=tmpWorkDir)
      except Errors.VibIOError as e:
         rmtree(tmpWorkDir)
         raise Errors.VibRecreateError(vibobj.id, "Cannot add payload: "
                                       "%s %s" % (e.msg, e.filename))

   try:
      vib.WriteVibFile(destfile)
   except Errors.VibIOError as e:
      raise Errors.VibRecreateError(vibobj.id, "Error while writing VIB file: "
                                    "%s" % e.msg)
   finally:
      rmtree(tmpWorkDir)

   verifyVibChecksumWithWorkaround()
   stageConfigSchema(arVib, destfile, tmpWorkDir)

def DepotFromImageProfile(imgprofile, depotdir, vibdownloadfn=VibDownloader,
                          vendor='Unknown', channels=[],
                          product=DEPOT_PRODUCT, versions=DEPOT_VERSIONS,
                          vendorcode='Unknown',
                          allowPartialDepot=False,
                          generateRollupBulletin=True,
                          vibDownloadArgs=None,
                          configSchemas=None):
   """Creates a complete depot from an image profile,
      including XML files, metadata.zip, and VIBs.

      Parameters:
         * imgprofile - An instance of ImageProfile, with the vibs attribute
                        containing a VibCollection with all of the VIBs in
                        vibIDs with the sourceurl attribute populated.
         * depotdir   - A directory to write out metadata.zip, XML files, and
                        VIB packages to.  The caller needs to create and
                        destroy this directory appropriately since this dir
                        will be needed until WriteBundleZip is called.
         * vibdownloadfn - A function taking params (destfile, vibobj) that is
                        responsible for downloading the Vib object vibobj to the
                        local path destfile.  If the original file is another
                        local path it should copy it.  This function could
                        create the VIBs from sources we have not anticipated
                        yet.
         * vendor     - String to use for the depot vendor name
         * channels   - A list of channel names to assign the depot to
         * product    - The string to use for productLineID
         * versions   - A list of the product versions supported by this depot
         * vendorcode - The vendor code.
         * allowPartialDepot - Whether allow export when VIB files missing.
                               When set, a partial depot, where some VIBs only
                               have metadata, may be created.
         * generateRollupBulletin - Whether generate the rollup bulletin for
                                    image profile.
         * vibDownloadArgs - Additional arguments passed to vibdownloadfn
         * configSchemas   - If not None, it provides the config schemas for
                             all VIBs (including reserved VIBs) that have
                             a config schema.
      Raises:
         VibDownloadError - if VIBs cannot be downloaded
         BundleIOError - error writing XML files to temp dir
         MetadataBuildError - error writing metadata.zip to temp dir
   """
   assert imgprofile.vibIDs.issubset(set(imgprofile.vibs.keys()))
   metazipbase = 'metadata.zip'

   if not HAVE_DOWNLOADER:
      raise ImportError("Failed to import downloader, offline bundle "
          "functionality not available.")

   allRelatedVibs = \
      VibCollection.VibCollection(imgprofile.vibs) + imgprofile.reservedVibs

   if configSchemas is None:
      # A temp dir to hold VIB config schemas until they are added as metadata.
      # It is used in host seeding case only.
      configSchemaDir = os.path.join(depotdir, 'configSchemas')
      if os.path.isdir(configSchemaDir):
         shutil.rmtree(configSchemaDir)
      os.makedirs(configSchemaDir)

      if vibDownloadArgs is None:
         vibDownloadArgs = dict()
      vibDownloadArgs['configSchemaDir'] = configSchemaDir

   hasDownloadIssue = False
   # Download all the VIBs
   for vibid in allRelatedVibs:
      localfn = os.path.join(depotdir, allRelatedVibs[vibid].GetRelativePath())
      try:
         vibdownloadfn(localfn, allRelatedVibs[vibid],
                       extraArgs=vibDownloadArgs)
      except EnvironmentError as e:
         raise Errors.VibDownloadError('', localfn, str(e))
      except Errors.VibDownloadError:
         hasDownloadIssue = True
         if not allowPartialDepot:
            raise

   ver = imgprofile.GetEsxVersion().split('-')[0]
   versions.append(ver)

   # Create the metadata.zip second
   meta = MetadataNode(url=metazipbase)
   if allowPartialDepot and hasDownloadIssue:
      meta.vibs = allRelatedVibs
   else:
      meta.vibs.FromDirectory(depotdir, ignoreinvalidfiles=True)

   meta.profiles.AddProfile(imgprofile)
   meta.bulletins += imgprofile.bulletins

   if generateRollupBulletin:
      bul = Bulletin.Bulletin(imgprofile.name, vendor=imgprofile.creator,
                              summary="Image Profile %s" % (imgprofile.name),
                              severity=Bulletin.Bulletin.SEVERITY_GENERAL,
                              urgency=Bulletin.Bulletin.URGENCY_MODERATE,
                              releasetype=Bulletin.Bulletin.RELEASE_ROLLUP,
                              category="Misc",
                              description=imgprofile.description,
                              platforms=[(ver, "", product) for ver in versions],
                              vibids=imgprofile.vibIDs,
                              )
      meta.bulletins.AddBulletin(bul)

   if imgprofile.baseimageID:
      meta.baseimages.AddBaseImage(imgprofile.baseimage)
   if imgprofile.addonID:
      meta.addons.AddAddon(imgprofile.addon)
   if imgprofile.manifestIDs:
      meta.manifests += imgprofile.manifests
   if imgprofile.solutionIDs:
      meta.solutions += imgprofile.solutions

   for bull in imgprofile.reservedComponents.GetComponents():
      meta.bulletins.AddBulletin(bull)

   if configSchemas is None:
      meta.configSchemas.FromDirectory(configSchemaDir)
      shutil.rmtree(configSchemaDir)
   else:
      meta.configSchemas = configSchemas

   for version in versions:
      meta.AddPlatform(product, version, channels=channels)
   meta.WriteMetadataZip(os.path.join(depotdir, metazipbase))

   # Create vendor-index.xml
   fn = 'vendor-index.xml'
   vidx = VendorIndex(name=vendor, code=vendorcode, indexfile=fn,
                            children=[meta])
   fpath = os.path.join(depotdir, fn)
   try:
      with open(fpath, 'wb') as vendorfile:
         vendorfile.write(vidx.ToString())
   except EnvironmentError as e:
      raise Errors.BundleIOError(fpath, "Error writing out vendor-index.xml "
                                 "for profile [%s]: %s" % (imgprofile.name,
                                                         str(e)))

   # create depot index.xml
   didx = DepotIndex(children=[vidx])
   fpath = os.path.join(depotdir, 'index.xml')
   try:
      with open(fpath, 'wb') as indexfile:
         indexfile.write(didx.ToString())
   except EnvironmentError as e:
      raise Errors.BundleIOError(fpath, "Error writing out index.xml for "
                                 "profile [%s]: %s" % (imgprofile.name,
                                                       str(e)))
