#!/usr/bin/python
########################################################################
# Copyright (C) 2010-2021 VMware, Inc.                                 #
# All Rights Reserved                                                  #
########################################################################

import os
import shutil
import tempfile

from .. import Database, Errors, Metadata, PayloadTar
from ..Utils import EsxGzip, Iso9660, SyslinuxConfig, XmlUtils
from ..Utils.HashedStream import HashedStream, HashError
from ..Utils.Misc import isString
from .ImageBuilder import ImageBuilder, createTgz, getSeekableFObj, resetFObj

etree = XmlUtils.FindElementTree()

"""This module provides a class for creating an ISO9660 image based on an
   image profile."""

class EsxIsoImage(ImageBuilder):
   "This class creates an ISO9660 image with the contents of an image profile."

   def __init__(self, imageprofile):
      """Class constructor.
            Parameters:
               * imageprofile - An instance of ImageProfile. The 'vibs'
                                attribute of the object must contain valid
                                references to the VIBs in the 'vibIDs'
                                property. Those references must include either
                                a file object or a valid remote location.
                                The bulletins/components contained in the
                                image profile must have their objects added to
                                the 'bulletins' attribute.
      """
      ImageBuilder.__init__(self, imageprofile)

   def _AddPayloads(self, volume, checkdigests=True, installer=True):
      # This method adds each payload to the ISO, while doing the right things
      # for various payload types.
      baseMiscTarfObj = tempfile.TemporaryFile()
      baseMiscTar = PayloadTar.PayloadTar(baseMiscTarfObj)

      imgpayldfobj = tempfile.TemporaryFile()
      imgpayldtar = PayloadTar.PayloadTar(imgpayldfobj)

      # Must generate unique names for certain payloads:
      self.imageprofile.GenerateVFATNames()
      for vibid in self.imageprofile.vibIDs:
         vibObj = self.imageprofile.vibs[vibid]
         for payload, fobj in vibObj.IterPayloads(checkDigests=checkdigests):
            if payload.payloadtype in payload.GZIP_TYPES:
               # We assume GenerateVFATNames() is working properly, and that
               # vibstates is populated.
               state = self.imageprofile.vibstates[vibid]
               payloadfn = state.payloads[payload.name]
            elif payload.payloadtype == payload.TYPE_TEXT:
               payloadfn = '-'.join([self.imageprofile.vibs[vibid].vendor,
                                     self.imageprofile.vibs[vibid].name,
                                     payload.name])
            else:
               payloadfn = payload.name

            if not payloadfn:
               msg = "VIB '%s' has payload with no name." % vibid
               raise Errors.VibFormatError(None, msg)

            # Some payloads are copied twice and must be seekable.
            fobj = getSeekableFObj(fobj)

            if payload.payloadtype in payload.NON_GZIP_TYPES:
               # Collect misc esx-base payloads.
               baseMiscTar.AddPayload(fobj, payload, payload.name)
               resetFObj(fobj)

            if payload.payloadtype in (payload.TYPE_VGZ, payload.TYPE_TGZ,
                                       payload.TYPE_BOOT_COM32_BIOS,
                                       payload.TYPE_TEXT):
               volume.AddFile(fobj, payloadfn, payload.size)
            elif (payload.payloadtype == payload.TYPE_INSTALLER_VGZ
                  and installer):
               volume.AddFile(fobj, payloadfn, payload.size)
            elif payload.payloadtype == payload.TYPE_BOOT:
               imgpayldtar.AddPayload(fobj, payload, payloadfn)
               resetFObj(fobj)
               volume.AddFile(fobj, payloadfn, payload.size)
            elif payload.payloadtype == payload.TYPE_UPGRADE and installer:
               volume.AddFile(fobj, "UPGRADE/" + payloadfn, payload.size)
            elif payload.payloadtype == payload.TYPE_BOOT_ISO_BIOS:
               record = volume.AddFile(fobj, payloadfn, payload.size)
               volume.SetBootImage(record)
            elif payload.payloadtype == payload.TYPE_ELTORITO_IMAGE_EFI:
               record = volume.AddFile(fobj, payloadfn, payload.size)
               volume.AddAltBootImage(record)
            elif payload.payloadtype == payload.TYPE_BOOT_LOADER_EFI:
               volume.AddFile(fobj, "EFI/BOOT/" + payloadfn, payload.size)

      imgpayldtar.close()
      imgpayldfobj.seek(0)

      baseMiscTar.close()
      baseMiscTarfObj.seek(0)
      volume.AddFile(baseMiscTarfObj, self.BASE_MISC_PAYLOADTAR_NAME)

      if installer:
         volume.AddFile(imgpayldfobj, self.PAYLOADTAR_NAME)

   def _AddReservedVibs(self, volume):
      """This method generates a tar file that contains reserved vibs.
         The tar file is added to the ISO as resvibs.tgz.
      """
      reservedVibTarObj = tempfile.TemporaryFile()
      super(EsxIsoImage, self)._AddReservedVibs(reservedVibTarObj)
      reservedVibTarObj.seek(0)
      volume.AddFile(reservedVibTarObj, self.RESERVED_VIBS_TAR_NAME)

   def _AddDatabase(self, volume):
      # This method generates a tar database from the image profile, writes it
      # to a temp file, then adds the temp file to the ISO image.
      db = Database.TarDatabase()
      db.PopulateWith(imgProfile=self.imageprofile)

      # check if vib signatures must be saved
      savesig = self.imageprofile.IsSecureBootReady()

      try:
         tmpf = tempfile.TemporaryFile()
         db.Save(dbpath=tmpf, savesig=savesig)
         tmpf.seek(0, 0)
      except Errors.EsxupdateError:
         # Preserve stack trace to enable troubleshooting the root cause.
         raise
      except EnvironmentError as e:
         msg = "Could not create temporary database: %s." % e
         raise Errors.DatabaseIOError(None, msg)

      volume.AddFile(tmpf, self.DATABASE_NAME)

   def _AddMetadataZip(self, volume):
      # This method generates a metadata.zip from the image profile.
      metadata = Metadata.Metadata()
      imgProfile = self.imageprofile
      metadata.profiles.AddProfile(imgProfile)
      metadata.bulletins += imgProfile.bulletins

      if imgProfile.baseimage:
         metadata.baseimages[imgProfile.baseimageID] = imgProfile.baseimage

      if imgProfile.addon:
         metadata.addons[imgProfile.addonID] = imgProfile.addon

      for vibid in self.imageprofile.vibIDs:
         metadata.vibs.AddVib(self.imageprofile.vibs[vibid])

      try:
         tmpdir = tempfile.mkdtemp()
      except IOError as e:
         msg = "Could not create temporary metadata directory: %s." % e
         raise Errors.MetadataIOError(msg)

      metapath = os.path.join(tmpdir, "metadata.zip")

      # Use nested try/except inside of try/finally to maintain Python 2.4
      # compatibility.
      try:
         try:
            metadata.WriteMetadataZip(metapath)
            # This is not efficient, but it saves us from having to worry about
            # keeping track of the temporary directory and cleaning it up
            # later.
            tmpf = tempfile.TemporaryFile()
            f = open(metapath, "rb")
            shutil.copyfileobj(f, tmpf)
            f.close()
            tmpf.seek(0)
         except IOError as e:
            msg = "Error copying metadata to temporary file: %s." % e
            raise Errors.MetadataIOError(msg)
      finally:
         try:
            shutil.rmtree(tmpdir)
         except Exception:
            pass

      volume.AddFile(tmpf, "upgrade/metadata.zip")

   def _AddProfileXml(self, volume):
      # Adds profile.xml to the upgrade directory on the ISO.
      profilexml = self.imageprofile.ToXml()
      xmltree = etree.ElementTree(profilexml)
      try:
         tmpf = tempfile.TemporaryFile()
         xmltree.write(tmpf)
         tmpf.seek(0)
      except Exception as e:
         # Specifically use Exception here, since we want to catch both I/O
         # errors, as well as XML serialization errors. Because we don't know
         # which specific etree implementation or what underlying parser it
         # will use, we also can't be sure of which exception we will get for
         # XML errors.
         msg = "Error writing temporary profile XML: %s." % e
         raise Errors.ProfileIOError(msg)

      volume.AddFile(tmpf, "upgrade/profile.xml")

   def _AddIsoLinuxConfig(self, volume, installer=True):
      # This method populates isolinux.cfg, using the profile name, and the
      # module order from the image profile. It writes the config to a temp
      # file, then adds the temp file to the ISO image.
      config = SyslinuxConfig.Config()
      config.default = "menu.c32"
      config.menutitle = "%s Boot Menu" % self.imageprofile.name
      config.timeout = 80
      config.nohalt = 1
      config.prompt = False

      label = config.AddLabel("install")
      if installer:
         label.menulabel = "%s ^Installer" % self.imageprofile.name
      else:
         label.menulabel = "%s ^System" % self.imageprofile.name
      label.kernel = "mboot.c32"
      label.append = "-c boot.cfg"

      label = config.AddLabel("hddboot")
      label.menulabel = "^Boot from local disk"
      label.localboot = "0x80"

      tmpf = tempfile.TemporaryFile()
      config.Write(tmpf)
      tmpf.seek(0, 0)

      volume.AddFile(tmpf, "isolinux.cfg")

   def _AddDiscinfo(self, volume):
      # This method populates .discinfo, a file which describes the contents
      # of the iso. The file format is not particualarly well-defined, so we have
      # a simple one with the product name and version
      # The final field, 'Virtual HW:' is used by Easy Install

      tmpf = tempfile.TemporaryFile()
      tmpf.write(("%(product)s\nVersion: %(version)s\n" % \
            {
               "product" : "ESXi",
               "version" : self.imageprofile.GetEsxVersion()
             }).encode())
      tmpf.seek(0, 0)

      volume.AddFile(tmpf, ".discinfo")

   def _AddEsxioDepot(self, volume, esxiodepot):
      """This method is for adding the ESXio-only depot to the ISO.
      """
      if isString(esxiodepot):
         try:
            depotfileobj = open(esxiodepot, 'rb')
         except Exception as e:
            raise Errors.FileIOError(esxiodepot, str(e))
      else:
         depotfileobj = esxiodepot

      try:
         depotTarObj = tempfile.TemporaryFile()
         createTgz(depotfileobj, 'esxio-depot.zip', depotTarObj)

         depotTarObj.seek(0)
         volume.AddFile(depotTarObj, self.ESXIO_DEPOT_TAR_NAME)
      except Exception as e:
         raise Errors.FileIOError(self.ESXIO_DEPOT_TAR_NAME, str(e))
      finally:
         if isString(esxiodepot):
            depotfileobj.close()

   def _AddBootCfg(self, volume, installer=True, kernelopts=None,
                   esxiodepot=None):
      # This method populates boot.cfg and efi/boot/boot.cfg.
      #
      # moduleroot is required for efi/boot/boot.cfg, but apparently doesn't hurt
      # anything for /boot.cfg.
      bootcfg = self._GetBootCfg(installer, moduleroot='/',
                                 kernelopts=kernelopts, isoImage=True,
                                 esxiodepot=esxiodepot)
      if not bootcfg:
         return

      tmpf = tempfile.TemporaryFile()
      bootcfg.write(tmpf)
      tmpf.seek(0, 0)

      volume.AddFile(tmpf, "boot.cfg")

      tmpf = tempfile.TemporaryFile()
      bootcfg.write(tmpf)
      tmpf.seek(0, 0)

      volume.AddFile(tmpf, "efi/boot/boot.cfg")

   def Write(self, f, checkdigests=True, insertmd5=True, installer=True,
             checkacceptance=True, kernelopts=None, esxiodepot=None):
      """Write out the ISO 9660 image to a file or file-like object.
            Parameters:
               * f               - A string giving a file name, or an object
                                   implementing the Python file protocol. The
                                   ISO image will be output to this path or
                                   file object.
               * checkdigests    - If True, payload digests will be verified
                                   when the ISO is written. Defaults to True.
               * insertmd5       - If True, an MD5 hash of the ISO contents
                                   will be inserted into the application data
                                   field of the ISO's primary volume
                                   descriptor. This is used by VUM Upgrade
                                   Manager for verifying the integrity of the
                                   image. Note that if this is True, the 'f'
                                   parameter must support rewinding the file
                                   pointer. Defaults to True.
               * installer       - Enable the installer in the booted image.
                                   Defaults to True.
               * checkacceptance - If True, validate the Acceptance Level of
                                   each VIB. If the validation fails, an
                                   exception is raised. Defaults to True.
               * kernelopts      - Additional kernel options as a dictionary.
               * esxiodepot      - File path or file object of the ESXio-only
                                   depot.
            Raises:
               * DatabaseIOError       - If unable to write the tar database to
                                         a temporary file.
               * ImageIOError          - If unable to write to a temporary file
                                         or the image output, or unable to
                                         compute the MD5 checksum of the image.
               * ProfileFormatError    - If the image profile has consistency
                                         errors.
               * VibDownloadError      - If unable to download one or more VIBs.
               * VibFormatError        - If one or more VIBs is not in proper
                                         VIB format.
               * VibIOError            - If unable to obtain the location of,
                                         or read data from, a VIB.
               * VibPayloadDigestError - If the calculated digest for one or
                                         more VIB payloads does not match the
                                         value given in VIB metadata.
               * FileIOError           - If the Esxio depot filename is invalid.
      """
      if isString(f):
         fobj = open(f, "w+b")
         def removeoutput():
            fobj.close()
            os.unlink(f)
      else:
         fobj = f
         def removeoutput():
            pass

      volume = Iso9660.Iso9660Volume()
      pvd = volume.primaryvolumedescriptor
      if insertmd5:
         # This is where we will store the md5 digest.
         pvd.applicationdata = "\0" * 16
      pvd.volumeid = self.imageprofile.name.upper()
      pvd.applicationid = "ESXIMAGE"

      try:
         self._CheckVibFiles(checkacceptance)
         self._AddPayloads(volume, checkdigests, installer)
         self._AddReservedVibs(volume)
         self._AddDatabase(volume)
         self._AddDiscinfo(volume)
         if installer:
            self._AddMetadataZip(volume)
            self._AddProfileXml(volume)
         if esxiodepot:
            self._AddEsxioDepot(volume, esxiodepot)
      except Exception:
         removeoutput()
         raise

      try:
         self._AddIsoLinuxConfig(volume, installer)
         self._AddBootCfg(volume, installer, kernelopts, esxiodepot)
      except IOError as e:
         removeoutput()
         msg = "Error writing boot configuration: %s." % e
         raise Errors.ImageIOError(str(fobj), msg)
      volume.Finalize()

      if insertmd5 and not hasattr(fobj, "seek") or not hasattr(fobj, "tell"):
         removeoutput() # not likely to be an actual file.
         msg = ("Can not insert MD5 digest into ISO image when writing to "
                "non-seekable output.")
         raise Errors.ImageIOError(str(fobj), msg)

      try:
         if insertmd5:
            # Application data is at offset 883 of primary volume descriptor:
            appdataoffset = fobj.tell() + (16 * 2048) + 883
            hashfobj = HashedStream(fobj, method="md5")
            volume.Write(hashfobj)
            isoend = fobj.tell()
            fobj.seek(appdataoffset)
            fobj.write(hashfobj.digest)
            fobj.seek(isoend)
         else:
            volume.Write(fobj)
      except (IOError, HashError, Iso9660.Iso9660Error) as e:
         removeoutput()
         msg = "Error occurred writing ISO image: %s." % e
         raise Errors.ImageIOError(str(fobj), msg)

      if isString(f):
         fobj.close()
