#!/usr/bin/python
########################################################################
# Copyright (C) 2010,2015-2018,2020-2021 VMWare, Inc.                  #
# All Rights Reserved                                                  #
########################################################################

import os
import shutil
import tempfile
import re

from .ImageBuilder import ImageBuilder, createTgz, getSeekableFObj, resetFObj
from .. import Database, Errors, Vib, PayloadTar
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"

   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):
      """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()
      for vibid in self.imageprofile.vibIDs:
         if not self.imageprofile.vibstates[vibid].boot:
            vib = self.imageprofile.vibs[vibid]
            _, _, 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 self.imageprofile.vibs[vibid].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]
            else:
               payloadfn = payload.name

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

            if payload.payloadtype in payload.NON_GZIP_TYPES:
               # 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 (payload.TYPE_VGZ, payload.TYPE_TGZ,
                                       payload.TYPE_BOOT_COM32_BIOS,
                                       payload.TYPE_BOOT_PXE_BIOS,
                                       payload.TYPE_BOOT_LOADER_EFI):
               if 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_INSTALLER_VGZ
                  and installer):
               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):
      bootcfg = self._GetBootCfg(installer, kernelopts=kernelopts,
                                 esxiodepot=esxiodepot)
      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):
      """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.
            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.
      """

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

   def WriteRecord(self, name, recordFile, pxeDir, treeMD5,
                   installer, targetType, opts=None, kernelopts=None,
                   esxiodepot=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.
               * treeMD5    - 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.
      """
      syslinuxCount = 0
      imgCount = 0

      localOpts = opts.copy() or {}

      localOpts["pxetype"] = targetType

      #
      # Add boot arguments string.
      #
      bootcfg = self._GetBootCfg(installer, kernelopts=kernelopts,
                                 esxiodepot=esxiodepot)
      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:
         for payload, fobj, in self.imageprofile.vibs[vibid].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.
      #
      payload_types = list(Vib.Payload.GZIP_TYPES)
      if not installer:
         payload_types.remove(Vib.Payload.TYPE_INSTALLER_VGZ)

      bootorder = self.imageprofile.GetBootOrder(payload_types)
      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" % (treeMD5,
                                           targetType,
                                           localOpts['bldtype'],
                                           key,
                                           localOpts[key])

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

