#!/usr/bin/python
########################################################################
# Copyright (c) 2010-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.
########################################################################

import os
import shutil
import tempfile

from .. import (BASE_VIBS, Database, Errors, Metadata,
                PayloadTar, Vib)
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
PAYLOAD_READ_CHUNKSIZE = 1024 * 1024
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."

   ESXIO_KS = 'esxio_ks.cfg'
   REGULAR_PAYLOAD_TYPES = Vib.Payload.TARDISK_TYPES + \
      (Vib.Payload.TYPE_BOOT_COM32_BIOS, Vib.Payload.TYPE_TEXT)

   def __init__(self, imageprofile, tmpDir=None):
      """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.
               * tmpDir       - Directory path for temporary file to be
                                created.
      """
      ImageBuilder.__init__(self, imageprofile)
      self.tmpDir = tmpDir

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

      imgpayldfobj = tempfile.TemporaryFile(dir=self.tmpDir)
      imgpayldtar = PayloadTar.PayloadTar(imgpayldfobj)

      # Must generate unique names for certain payloads:
      self.imageprofile.GenerateVFATNames(platform=platform)
      for vibid in self.imageprofile.vibIDs:
         vibObj = self.imageprofile.vibs[vibid]
         if platform and not vibObj.HasPlatform(platform):
            continue
         for payload, fobj in vibObj.IterPayloads(checkDigests=checkdigests):
            if payload.payloadtype in payload.ALL_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, tmpDir=self.tmpDir)

            # 'useropts' is not nongzip but it should be stored in basemisc.tgz
            # for host seeding to re-create esx-base VIB with the original
            # contents that might be non-empty
            if (payload.payloadtype in payload.NON_GZIP_TYPES or
                (vibObj.name in BASE_VIBS and payload.name in
                 PayloadTar.PayloadTar.ADDITIONAL_BASEMISC_PAYLOADS)):
               # Collect misc esx-base payloads.
               baseMiscTar.AddPayload(fobj, payload, payload.name)
               resetFObj(fobj)

            if payload.payloadtype in self.REGULAR_PAYLOAD_TYPES:
               if (payload.payloadtype == payload.TYPE_INSTALLER_VGZ
                     and not installer):
                  continue
               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, platform=None, createEmptyTar=False):
      """This method generates a tar file that contains reserved vibs,
         filter by platform if specified.
         The tar file is added to the ISO as resvibs.tgz.
         createEmptyTar - If True, only add an empty resvibs.tgz and skip all
                          reserved VIBs. Defaults to False.
      """
      reservedVibTarObj = tempfile.TemporaryFile(dir=self.tmpDir)
      super(EsxIsoImage, self)._AddReservedVibs(
                              reservedVibTarObj, platform=platform,
                              createEmptyTar=createEmptyTar)
      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(dir=self.tmpDir)
         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])

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

      for comp in self.imageprofile.reservedComponents.GetComponents():
         metadata.bulletins.AddBulletin(comp)

      try:
         tmpdir = tempfile.mkdtemp(dir=self.tmpDir)
      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(dir=self.tmpDir)
            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(dir=self.tmpDir)
         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(dir=self.tmpDir)
      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(dir=self.tmpDir)
      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(dir=self.tmpDir)
         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, esxioKsfile=None, platform=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.
      """
      if esxioKsfile:
         if kernelopts is None:
            kernelopts = dict()
         kernelopts['ks'] = 'file://' + self.ESXIO_KS

      bootcfg = self._GetBootCfg(installer, moduleroot='/',
                                 kernelopts=kernelopts, isoImage=True,
                                 esxiodepot=esxiodepot, esxioKsfile=esxioKsfile,
                                 platform=platform)
      if not bootcfg:
         return

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

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

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

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

   def _AddKsFile(self, volume, esxioKsfile):
      """This method is for adding the kickstart file to the ISO.
      """

      try:
         ksfileobj = open(esxioKsfile, 'rb')
      except Exception as e:
         raise Errors.FileIOError(esxioKsfile, str(e))

      try:
         ksTarObj = tempfile.TemporaryFile(dir=self.tmpDir)
         createTgz(ksfileobj, self.ESXIO_KS, ksTarObj)

         ksTarObj.seek(0)
         volume.AddFile(ksTarObj, self.ESXIO_KS_TAR_NAME)
      except Exception as e:
         raise Errors.FileIOError(self.ESXIO_KS_TAR_NAME, str(e))
      finally:
         ksfileobj.close()

   def Write(self, f, checkdigests=True, insertHash=True, installer=True,
             checkacceptance=True, kernelopts=None, esxiodepot=None,
             esxioKsfile=None, platform=None, partialDepot=False,
             storeReservedVibs=True):
      """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.
               * insertHash      - If True, insert a MD5, version and a SHA256
                                   hash of the ISO contents 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.
               * esxioKsfile     - Kickstart file to be included for scripted
                                   ESXio install.
               * platform        - SoftwarePlatform productLineID whose VIB
                                   payloads are to be written to the ISO. VIBs
                                   for other platforms are ignored but metadata
                                   is still included.
               * partialDepot    - Flag to specify if the depot which was used
                                   to create the image profile is a partial
                                   depot or not. If set to True, the platform
                                   arg must also be provided.
               * storeReservedVibs-If True, all reserved VIBs are added to the
                                   tar file. If False, add an empty
                                   resvibs.tgz. Defaults to True.

            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/SHA256 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 or
                                         kickstart file 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 insertHash:
         # We will store MD5 (16 bytes) + version (1 byte) + SHA256 (32 bytes).
         # For backward-compatibility, MD5 will occupy the first 16 bytes.
         # Version 0x20: legacy ISO, MD5 at 0-15
         # Version 0x21: ISO with both MD5 at 0-15 and SHA256 at 17-48.
         pvd.applicationdata = "\0" * (16 + 1 + 32)
      pvd.volumeid = self.imageprofile.name.upper()
      pvd.applicationid = "ESXIMAGE"

      try:
         self._CheckVibFiles(checkacceptance, platform, partialDepot)
         self._AddPayloads(volume, checkdigests, installer, platform)
         self._AddReservedVibs(volume, platform=platform,
                               createEmptyTar=not storeReservedVibs)
         self._AddDatabase(volume)
         self._AddDiscinfo(volume)
         if installer:
            self._AddMetadataZip(volume)
            self._AddProfileXml(volume)
         if esxiodepot:
            self._AddEsxioDepot(volume, esxiodepot)
         if esxioKsfile:
            self._AddKsFile(volume, esxioKsfile)
      except Exception:
         removeoutput()
         raise

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

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

      try:
         if insertHash:
            # Application data is at offset 883 of primary volume descriptor:
            # Insert MD5/SHA256 hash. SHA256 is calculated with all 49 bytes of
            # application data zeroed, while for backward compatibility, MD5 is
            # calculated with SHA256 and version already inserted and treated
            # like regular ISO contents.
            appdataoffset = fobj.tell() + (16 * 2048) + 883
            hashfobjsha256 = HashedStream(fobj, method="sha256")
            volume.Write(hashfobjsha256)
            isoend = fobj.tell()

            # Insert version and 32 bytes SHA256 hash into bytes (17 to 48).
            # To maintain backward compatibility, SHA256 hash is calculated and
            # inserted first.
            fobj.seek(appdataoffset + 16)
            fobj.write(b"\x21")
            fobj.write(hashfobjsha256.digest)

            # Reset the file object and calculate the MD5 hash by reading the
            # ISO with SHA256 filled, then insert MD5 hash before SHA256.
            # The verification sequence should reverse (MD5 and then SHA256)
            # to maintain backward compatibility.
            fobj.seek(0)
            hashfobjMD5 = HashedStream(fobj, method="md5")
            data = hashfobjMD5.read(PAYLOAD_READ_CHUNKSIZE)
            while data:
               data = hashfobjMD5.read(PAYLOAD_READ_CHUNKSIZE)

            # Insert 16 bytes MD5 hash into bytes (0 to 15).
            fobj.seek(appdataoffset)
            fobj.write(hashfobjMD5.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()
