#!/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
import re

from .ImageBuilder import ImageBuilder, createTgz, getSeekableFObj, resetFObj
from .. import BASE_VIBS, Database, Errors, PayloadTar, Vib
from ..Utils import EsxGzip
from ..Utils.Misc import isString

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

class EsxPxeImage(ImageBuilder):
   "This class creates a PXE image with the contents of an image profile."

   DATABASE_NAME = "imgdb.tgz"
   PAYLOADTAR_NAME = "imgpayld.tgz"
   REGULAR_PAYLOAD_TYPES = Vib.Payload.TARDISK_TYPES + \
      (Vib.Payload.TYPE_BOOT_COM32_BIOS, Vib.Payload.TYPE_BOOT_PXE_BIOS,
       Vib.Payload.TYPE_BOOT_LOADER_EFI)

   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)

   @staticmethod
   def _CopyFileObjToFileName(srcfobj, destfpath, length=None):
      destdir = os.path.dirname(destfpath)
      if not os.path.exists(destdir):
         os.makedirs(destdir)

      with open(destfpath, 'wb') as newfobj:
         if length:
            shutil.copyfileobj(srcfobj, newfobj, length)
         else:
            shutil.copyfileobj(srcfobj, newfobj)

   def _DeployVib(self, pxedir, checkdigests=True, installer=True,
                  platform=None):
      """Deploy all VIBs to the PXE directory.

      Boot-VIBs are extracted to the PXEDIR directory. Other extra VIBs are
      copied to PXEDIR/vibs/.
      """
      baseMiscTarPath = os.path.join(pxedir, self.BASE_MISC_PAYLOADTAR_NAME)
      baseMiscTar = PayloadTar.PayloadTar(baseMiscTarPath)

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

      # Must generate unique names for certain payloads:
      self.imageprofile.GenerateVFATNames(platform=platform)
      for vibid in self.imageprofile.vibIDs:
         vib = self.imageprofile.vibs[vibid]
         if platform and not vib.HasPlatform(platform):
            continue
         if not self.imageprofile.vibstates[vibid].boot:
            _, _, name, version = vibid.split('_')
            vibName = "%s-%s.i386.vib" % (name, version)
            vibDir = os.path.join(pxedir, 'vibs')
            if not os.path.exists(vibDir):
               os.makedirs(vibDir)
            vib.WriteVibFile(os.path.join(vibDir, vibName))
            continue

         for payload, fobj in vib.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]
            else:
               payloadfn = payload.name

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

            # '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
                (vib.name in BASE_VIBS and payload.name in
                 PayloadTar.PayloadTar.ADDITIONAL_BASEMISC_PAYLOADS)):
               # Collect misc esx-base payloads, many of them are copied
               # twice and must be seekable.
               fobj = getSeekableFObj(fobj)
               baseMiscTar.AddPayload(fobj, payload, payload.name)
               resetFObj(fobj)

            if payload.payloadtype in self.REGULAR_PAYLOAD_TYPES:
               if payload.payloadtype == payload.TYPE_INSTALLER_VGZ:
                  if not installer:
                     continue
               elif re.match('boot.*\.efi$', payloadfn, re.IGNORECASE):
                  payloadfn = 'mboot.efi'
               newfpath = os.path.join(pxedir, payloadfn)
               self._CopyFileObjToFileName(fobj, newfpath, payload.size)
            elif payload.payloadtype == payload.TYPE_BOOT:
               newfpath = os.path.join(pxedir, payloadfn)
               self._CopyFileObjToFileName(fobj, newfpath, payload.size)
               with open(newfpath, 'rb') as newfobj:
                  imgpayldtar.AddPayload(newfobj, payload, payloadfn)

      baseMiscTar.close()

      imgpayldtar.close()
      imgpayldfobj.seek(0)

      if installer:
         newfpath = os.path.join(pxedir, self.PAYLOADTAR_NAME)
         self._CopyFileObjToFileName(imgpayldfobj, newfpath)
         imgpayldfobj.close()

   def _AddReservedVibs(self, pxedir):
      """This method generates a tar file that contains reserved vibs.
         The tar file is added to the PXE dir as resvibs.tgz.
      """
      reservedVibTarPath = os.path.join(pxedir, self.RESERVED_VIBS_TAR_NAME)
      super(EsxPxeImage, self)._AddReservedVibs(reservedVibTarPath)

   def _AddDatabase(self, pxedir):
      # This method generates a tar database from the image profile, writes it
      # to a temp file, then adds the temp file to the PXE directory
      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)

      newfpath = os.path.join(pxedir, self.DATABASE_NAME)
      self._CopyFileObjToFileName(tmpf, newfpath)

   def _AddBootCfg(self, pxedir, installer=True, kernelopts=None, prefix=None,
                   esxiodepot=None, platform=None):
      bootcfg = self._GetBootCfg(installer, kernelopts=kernelopts,
                                 esxiodepot=esxiodepot, platform=platform)
      assert bootcfg is not None,\
            "No module in image profile '%s'..." % (self.imageprofile.name)
      if prefix:
         bootcfg.prefix = prefix
      bootcfg.write(os.path.join(pxedir, 'boot.cfg'))

   def _WriteIpxeConf(self, pxeDir, pxeUrl):
      """ Write an iPXE boot script that chainloads mboot"""
      ipxe = os.path.join(pxeDir, 'ipxe.conf')
      if not pxeUrl.endswith("/"):
         pxeUrl = pxeUrl + "/"
      conf = ['#!ipxe',
              'iseq ${platform} efi && goto efi ||',
              'chain %smboot.c32 -S1 -c %sboot.cfg'
              ' BOOTIF=01-${netX/mac:hexhyp}' % (pxeUrl, pxeUrl),
              'exit',
              ':efi',
              'chain %smboot.efi -S1 -c %sboot.cfg' % (pxeUrl, pxeUrl)]
      with open(ipxe, 'w') as f:
         f.write('\n'.join(conf))

   def _AddEsxioDepot(self, pxeDir, esxiodepot):
      """ Write the ESXio-only depot to the PXE directory.
      """

      if isString(esxiodepot):
         try:
            depotfileobj = open(esxiodepot, 'rb')
         except Exception as e:
            raise Errors.FileIOError(esxiodepot, str(e))
      else:
         depotfileobj = esxiodepot

      try:
         depotTarPath = os.path.join(pxeDir, self.ESXIO_DEPOT_TAR_NAME)
         createTgz(depotfileobj, "esxio-depot.zip", depotTarPath)
      except Exception as e:
         raise Errors.FileIOError(self.ESXIO_DEPOT_TAR_NAME, str(e))
      finally:
         if isString(esxiodepot):
            depotfileobj.close()

   def Write(self, pxedir, pxeUrl=None, checkdigests=True, installer=True,
             checkacceptance=True, kernelopts=None, bootCfgPrefix=None,
             esxiodepot=None, platform=None):
      """Write out the files to a PXE directory.
            Parameters:
               * pxedir          - A string giving the absolute path to a
                                   directory.  Files for the PXE will be written
                                   to this directory.
               * pxeUrl          - A string providing the url http location
                                   of PXE output directory. This will be used
                                   to generate ipxe.conf.
               * checkdigests    - If True, payload digests will be verified
                                   when the PXE is written. 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 other than
                                   feature states, as a dictionary.
               * bootCfgPrefix   - The prefix to use in the boot configuration
                                   file.
               * esxiodepot      - File path or file object of the ESXio-only
                                   depot.
               * platform        - If set, payloads of VIBs for other platforms
                                   are skipped, but metadata of them will be
                                   included.
            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 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 is invalid.
      """

      self._CheckVibFiles(checkacceptance, platform)
      self._DeployVib(pxedir, checkdigests, installer, platform)
      self._AddReservedVibs(pxedir)
      self._AddDatabase(pxedir)
      if esxiodepot:
         self._AddEsxioDepot(pxedir, esxiodepot)
      self._AddBootCfg(pxedir, installer, kernelopts=kernelopts,
                       prefix=bootCfgPrefix, esxiodepot=esxiodepot,
                       platform=platform)
      if pxeUrl is not None:
         self._WriteIpxeConf(pxedir, pxeUrl)

   def WriteRecord(self, name, recordFile, pxeDir, treeSHA256,
                   installer, targetType, opts=None, kernelopts=None,
                   esxiodepot=None, platform=None):
      """Write out a PXE record file for use by the pxe-boot perl script.
            Parameters:
               * name       - A name for the PXE image.
               * recordFile - The full path to the PXE record file that we wish
                              to write to.
               * pxeDir     - The full path to the directory that contains the
                              staged PXE files.
               * treeSHA256    - An hashsum (of the path to your tree) that's
                              used to distinguish between your different trees.
               * installer  - Enables the installer in the PXE image.
               * targetType - The build type that we're using (obj, beta, release)
               * opts       - Any additional options that need to be passed to
                              the pxe-boot script.
               * kernelopts - Additional kernel options other than
                              feature states, as a dictionary.
               * esxiodepot - File path or file object of the ESXio-only depot.
               * platform   - SoftwarePlatform productLineID whose VIB payloads
                              are to be written into boot.cfg.
      """
      syslinuxCount = 0
      imgCount = 0

      localOpts = opts.copy() or {}

      localOpts["pxetype"] = targetType

      #
      # Add boot arguments string.
      #
      bootcfg = self._GetBootCfg(installer, kernelopts=kernelopts,
                                 esxiodepot=esxiodepot, platform=platform)
      assert bootcfg is not None,\
            "No module in image profile '%s'..." % (self.imageprofile.name)
      localOpts["bootargs"] = bootcfg.kerneloptToStr(bootcfg.kernelopt)

      #
      # Add options of the form "syslinux.0", "syslinux.1"... to define syslinux
      # modules. Specify these in the config file as paths relative to the
      # config file we'll be writing.
      #

      # Internal deployed PXE will contain extra syslinux modules, add
      # them before scanning the vibs.
      pxeBootFiles = []
      if localOpts['arch'] == 'x64':
         pxeBootFiles += ['gpxelinux.0', 'ifgpxe.c32', 'ipxe-undionly.0']
      for name in pxeBootFiles:
         localOpts['syslinux.%s' % syslinuxCount] = name
         syslinuxCount += 1

      for vibid in self.imageprofile.vibIDs:
         vib = self.imageprofile.vibs[vibid]
         if platform and not vib.HasPlatform(platform):
            continue
         for payload, fobj, in vib.IterPayloads():
            if payload.payloadtype in [payload.TYPE_BOOT_COM32_BIOS,
                                       payload.TYPE_BOOT_PXE_BIOS,
                                       payload.TYPE_BOOT_LOADER_EFI]:
               if re.match('boot.*\.efi$', payload.name, re.IGNORECASE):
                  payload.name = 'mboot.efi'
               filePath = os.path.join(pxeDir, payload.name)
               relPath = os.path.relpath(filePath,
                                         os.path.dirname(recordFile))
               localOpts["syslinux.%s" % syslinuxCount] = relPath
               syslinuxCount += 1

      #
      # Add options of the form "image.0", "image.1".... to define
      # modules to be loaded.  Specify these in the config file as
      # paths relative to the config file we'll be writing.
      #
      payloadTypes = list(Vib.Payload.ALL_GZIP_TYPES)
      if not installer:
         payloadTypes.remove(Vib.Payload.TYPE_INSTALLER_VGZ)

      bootorder = self.imageprofile.GetBootOrder(payloadTypes,
                                                 platform=platform)
      modules = [ p.localname for (vibid, p) in bootorder ]

      modules.append(self.DATABASE_NAME)
      modules.append(self.BASE_MISC_PAYLOADTAR_NAME)
      modules.append(self.RESERVED_VIBS_TAR_NAME)
      if esxiodepot:
         modules.append(self.ESXIO_DEPOT_TAR_NAME)
      if installer:
         modules.append(self.PAYLOADTAR_NAME)

      for m in modules:
         filePath = os.path.join(pxeDir, m)
         relPath = os.path.relpath(filePath,
                                   os.path.dirname(recordFile))
         localOpts["image.%s" % imgCount] = relPath
         imgCount += 1

      # Full version of the build
      localOpts["esxVersion"] = self.imageprofile.GetEsxVersion()

      #
      # The format for these files is:
      # <treeID>.<targetType>.<buildType>.key = val
      #
      output = ''
      for key in localOpts:
         output += "%s.%s.%s.%s = %s\n" % (treeSHA256,
                                           targetType,
                                           localOpts['bldtype'],
                                           key,
                                           localOpts[key])

      with open(recordFile, 'w') as record:
         record.write(output)

