#!/usr/bin/python

'''
Copyright 2018-2023 VMware, Inc.
All rights reserved. -- VMware Confidential
'''

'''
This module generates an ESX Raw Disk Image based on an Image Profile. This
raw image can later be dd'ed onto a SD card or USB stick.
'''

import datetime
import errno
import os
import shutil
from tempfile import TemporaryDirectory

from .. import BASE_VIBS, Database, Depot, Errors, PayloadTar, Vib
from .ImageBuilder import ImageBuilder, getSeekableFObj, resetFObj
from ..OfflineBundle import OfflineBundle
from ..Utils import XmlUtils, Misc
from ..Utils.BootCfg import BootCfg

from systemStorage.ddImage import DDImage
from systemStorage.vfat import Libvfat


class EsxRawImage(ImageBuilder):
   """This class creates a Raw Disk Image based on an image profile"""

   def __init__(self, imageProfile, mmdBin=None, mcopyBin=None, vfatSo=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.

               * mcopyBin:      Path the mcopy executable
               * mmdBin:        Path to the mmd executable
               * vfatSo:        Path to the vfat shared library
      """
      ImageBuilder.__init__(self, imageProfile)

      self.mcopyBin = mcopyBin
      self.mmdBin = mmdBin

      if vfatSo is not None:
         Libvfat.path = vfatSo

   def Write(self, imgFilename, checkdigests=True, checkacceptance=True,
             kernelopts=None, platform=None, extraBootMods=None, stagedir=None):
      """Writes a raw ESX image to a file

         All bootloader, bootbank and locker files are extracted from
         an image profile and placed in their respective staging
         directories. These directories are then passed into DDImage
         which performs the actual write.

         Parameters:
            * imgFilename: The filename of the image to write
            * checkDigests: Set to True to check VIB digests
            * checkacceptance: Set to True to check VIB acceptance
            * kernelopts: dictionary of kernel options for this image
            * platform: If set, then skip VIBs not meant for the platform.
            * extraBootMods: Optional list of additional boot modules
            * stagedir: The location where staging happens. Defauts to /tmp.
      """
      if stagedir and not os.path.isdir(stagedir):
         raise FileNotFoundError(errno.ENOTDIR, os.strerror(errno.ENOTDIR), stagedir)
      with TemporaryDirectory(dir=stagedir) as stageDir:

         bootloaderDir = os.path.join(stageDir, 'bootloader')
         os.mkdir(bootloaderDir)
         self._GetBootloaderFiles(bootloaderDir, checkdigests, checkacceptance,
                                  platform)

         bootbank1Dir = os.path.join(stageDir, 'bootbank1Dir')
         os.mkdir(bootbank1Dir)
         bootbank2Dir = os.path.join(stageDir, 'bootbank2Dir')
         os.mkdir(bootbank2Dir)
         self._GetBootbankFiles(bootbank1Dir, bootbank2Dir, kernelopts,
                                checkdigests, checkacceptance,
                                platform=platform, extraBootMods=extraBootMods)

         lockerDir = os.path.join(stageDir, 'lockerDir')
         os.mkdir(lockerDir)
         self._GetLockerFiles(lockerDir)

         # If the image profile has reserved vibs, copy them to
         # stageDir/hostSeedDir/reservedVibs
         # hostSeedDir is temporary location, reservedVibs gets
         # copied to bootbank2.
         hostSeedDir = os.path.join(stageDir, 'hostSeedDir')
         os.makedirs(os.path.join(hostSeedDir, 'reservedVibs'))
         os.makedirs(os.path.join(hostSeedDir, 'esxioCachedVibs'))
         self._AddReservedVibs(os.path.join(hostSeedDir, 'reservedVibs'))
         self._CacheEsxioVibs(os.path.join(hostSeedDir, 'esxioCachedVibs'))

         ddImage = DDImage(imgFilename, mmdBin=self.mmdBin,
                           mcopyBin=self.mcopyBin)
         ddImage.write(bootloaderDir, bootbank1Dir, lockerDir, hostSeedDir)

   def _GetBootbankFiles(self, bootbank1Dir, bootbank2Dir, kernelopts,
                         checkdigests, checkacceptance, platform=None,
                         extraBootMods=None):
      """Extract the bootbank payloads into the bootbank directories

         Parameters:
            * bootbank1Dir: The directory to place bootbank1 files
            * bootbank2Dir: The directory to place bootbank2 files
            * kernelopts: A dictionary of kernel options for this image
            * checkDigests: Set to True to check Vib digests
            * checkacceptance: Set to True to check VIB acceptance
            * platform: If set, then skip VIBs not meant for the platform.
            * extraBootMods: Optional list of additional boot modules
      """
      self._GetBootbankVibs(bootbank1Dir, checkdigests, checkacceptance,
                            platform)

      # Stage extra boot modules into bootbank.
      if extraBootMods is not None:
         for mod in extraBootMods:
            dest = os.path.join(bootbank1Dir, os.path.basename(mod))
            shutil.copyfile(mod, dest)

      self._CreateBootCfg(bootbank1Dir, True, kernelopts, platform=platform,
                          extraBootMods=extraBootMods)
      self._CreateDatabase(bootbank1Dir)

      self._CreateBootCfg(bootbank2Dir, False, kernelopts, platform=platform)

   def _GetBootloaderFiles(self, bootloaderDir, checkdigests, checkacceptance,
                           platform=None):
      """Extract the bootloader payloads into a bootloader directory

         Parameters:
            * bootloaderDir: The directory to place the bootloader files
            * checkDigests: Set to True to check VIB digests
            * checkacceptance: Set to True to check VIB acceptance
            * kernelopts: dictionary of kernel options for this image
            * platform: If set, then skip VIBs not meant for the platform.
      """
      self._CheckVibFiles(checkacceptance=checkacceptance)

      for vibid in self.imageprofile.vibIDs:
         vibObj = self.imageprofile.vibs[vibid]
         if vibObj.vibtype != Vib.BaseVib.TYPE_BOOTBANK:
            continue

         if platform and not vibObj.HasPlatform(platform):
            continue

         for payload, fobj in vibObj.IterPayloads(checkDigests=checkdigests):
            if payload.payloadtype not in [Vib.Payload.TYPE_BOOT_COM32_BIOS,
                                           Vib.Payload.TYPE_BOOT_LOADER_EFI,
                                           Vib.Payload.TYPE_UPGRADE]:
               continue

            dest = os.path.join(bootloaderDir, payload.name)
            with open(dest, 'wb') as newfobj:
               shutil.copyfileobj(fobj, newfobj, length=payload.size)

      # Rename the bootloader files which have different names in the base
      # VIB and inside the 'btldr' tardisk. We want the latter.
      fixups = {'BOOTAA64.EFI': 'mboot64.efi',
                'BOOTx64.EFI': 'mboot64.efi',
                'safeboot.efi': 'safeboot64.efi'}

      for basename in os.listdir(bootloaderDir):
         if basename in fixups:
            os.rename(os.path.join(bootloaderDir, basename),
                      os.path.join(bootloaderDir, fixups[basename]))

   def _GetBootbankVibs(self, stageDir, checkdigests, checkacceptance,
                        platform=None):
      """Stages Bootbank Vib payloads from the image profile, and packages
         miscellaneous esx-base payloads in a tar.

         Parameters:
            * stageDir: Staging Dir to copy the bootbank vibs
            * checkDigests: Set to True to check Vib digests
            * checkacceptance: Set to True to check VIB acceptance
            * platform: If set, then skip VIBs not meant for the platform.

      """
      installTime = datetime.datetime.now(XmlUtils._utctzinfo)
      self._CheckVibFiles(checkacceptance=checkacceptance)

      payloadTypes = [Vib.Payload.TYPE_TGZ, Vib.Payload.TYPE_VGZ,
                      Vib.Payload.TYPE_BOOT]

      self.imageprofile.GenerateVFATNames(platform=platform)

      baseMiscTarPath = os.path.join(stageDir, self.BASE_MISC_PAYLOADTAR_NAME)
      baseMiscTar = PayloadTar.PayloadTar(baseMiscTarPath)

      for vibid in self.imageprofile.vibIDs:
         vibObj = self.imageprofile.vibs[vibid]
         # Add install date for all VIBs as at least locker VIB is assumed
         # to have one.
         vibObj.installdate = installTime
         if platform and not vibObj.HasPlatform(platform):
            continue
         if vibObj.vibtype == Vib.BaseVib.TYPE_BOOTBANK:
            for payload, fobj in vibObj.IterPayloads(checkDigests=checkdigests):
               # '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. A VIB in zip being extracted
                  # is not seekable, we have to write to a temporary file first
                  # in order to read the payload twice.
                  fobj = getSeekableFObj(fobj)
                  baseMiscTar.AddPayload(fobj, payload, payload.name)
                  resetFObj(fobj)

               if payload.payloadtype in payloadTypes:
                  state = self.imageprofile.vibstates[vibid]
                  payloadFilename = state.payloads[payload.name]

                  vibFilename = os.path.join(stageDir, payloadFilename)
                  with open(vibFilename, 'wb') as newfobj:
                     shutil.copyfileobj(fobj, newfobj, length=payload.size)
               fobj.close()

      baseMiscTar.close()

   def _AddReservedVibs(self, dstDir):
      """This method downloads the all the reserved vibs in the image profile
         (except tools vib) into the destination directory."""
      try:
         for vibid in self.imageprofile.reservedVibIDs:
            vib = self.imageprofile.reservedVibs[vibid]
            # locker VIB is specially reserved for cases where
            # bootbank and locker are not synced, a full copy is not needed.
            if vib.vibtype == vib.TYPE_LOCKER or not vib.relativepath:
               continue

            dstFile = os.path.join(dstDir, '%s.vib' %vibid)
            Depot.VibDownloader(dstFile, vib)
      except Exception as e:
         msg = "Could not download and package reserved VIBs. %s" % str(e)
         raise Errors.VibDownloadError(None, None, msg)

   def _CacheEsxioVibs(self, dstDir):
      """Downloads the all the ESXio vibs in the image profile into the
         destination directory.
      """
      if len(self.imageprofile.GetSoftwarePlatforms(fillDefaultValue=False,
         baseEsxOnly=True)) == 1:
         # If image profile is not unified, then no need to cache.
         return

      try:
         for vibid in self.imageprofile.vibIDs:
            vibObj = self.imageprofile.vibs[vibid]
            if vibObj.HasPlatform(Vib.SoftwarePlatform.PRODUCT_EMBEDDEDESX) or \
               not vibObj.HasPlatform(Vib.SoftwarePlatform.PRODUCT_ESXIO_ARM):
               continue
            dstFile = os.path.join(dstDir, '%s.vib' % vibid)
            Depot.VibDownloader(dstFile, vibObj)
      except Exception as e:
         msg = "Could not download and package ESXio VIBs: %s" % str(e)
         raise Errors.VibDownloadError(None, None, msg)

   def _CreateBootCfg(self, bootCfgDir, isPrimaryBootbank, kernelopts,
                      platform=None, extraBootMods=None):
      """Creates a bootbank configuration file

         Parameters:
            * bootCfgDir: Staging Dir to create the boot configuration file.
            * isPrimaryBootbank: When set to true this will create a
                                 bootbank config contain a kernel image
                                 and a list of modules. When set to false
                                 an empty bootbank is created.
            * kernelopts: A dictionary of kernel options for this image
            * platform: If set, then skip VIBs not meant for the platform.
      """
      if isPrimaryBootbank:
         if kernelopts is not None:
            kernelopts.update({'esxDDImage': 'TRUE'})
         else:
            kernelopts = {'esxDDImage': 'TRUE'}

      bootcfg = self._GetBootCfg(installer=False,
                                 kernelopts=kernelopts,
                                 bootbankVibOnly=isPrimaryBootbank,
                                 appendResVibsTgz=False,
                                 platform=platform,
                                 extraBootMods=extraBootMods)

      assert bootcfg is not None,\
         "No module in image profile '%s'..." % (self.imageprofile.name)

      bootcfg.updated = 1
      if isPrimaryBootbank:
         bootcfg.bootstate = BootCfg.BOOTSTATE_SUCCESS
      else:
         bootcfg.kernel = ''
         bootcfg.modules = []
         bootcfg.bootstate = BootCfg.BOOTSTATE_EMPTY

      bootCfgFilename = os.path.join(bootCfgDir, 'boot.cfg')
      f = open(bootCfgFilename, 'wb')
      bootcfg.write(f)
      f.close()

   def _CreateDatabase(self, databaseDir):
      """ This method generates a tar database from the image profile and vibs
          and writes it to a file.

         Parameters:
            * databaseDir: Dir to create the database
      """
      profile = self.imageprofile.Copy()
      db = Database.TarDatabase()
      db.PopulateWith(imgProfile=profile)

      for vibid in self.imageprofile.vibIDs:
         if self.imageprofile.vibs[vibid].vibtype == Vib.BaseVib.TYPE_LOCKER:
            # Locker VIB is loaded from locker and will not be in the
            # bootbank DB.
            db.profile.RemoveVib(vibid)
            db.vibs.RemoveVib(vibid)
      try:
         # check if vib signatures must be saved
         savesig = self.imageprofile.IsSecureBootReady()
         databaseFilename = os.path.join(databaseDir, self.DATABASE_NAME)
         database = open(databaseFilename, 'wb')
         db.Save(dbpath=database, savesig=savesig)
         database.close()
      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)

   def _GetLockerFiles(self, lockerDir):
      """This method untars the tools payload into the a directory
         and create the locker database.
         Both the extracted tools files and the db will be in
         lockerDir/packages as LockerInstaller anticipates all locker files
         to be in /locker/packages.

         Parameters:
            * lockerDir: Locker Dir to create the database
      """
      pkgDir = os.path.join(lockerDir, 'packages')
      dbDir = os.path.join(pkgDir, 'var', 'db', 'locker')

      db = Database.Database(dbDir, addprofile=False)
      for vibid in self.imageprofile.vibIDs:
         if self.imageprofile.vibs[vibid].vibtype == Vib.BaseVib.TYPE_LOCKER:
            db.vibs.AddVib(self.imageprofile.vibs[vibid])
            vibObj = self.imageprofile.vibs[vibid]
            for payload, fobj in vibObj.IterPayloads(checkDigests=True):
               if payload.payloadtype == payload.TYPE_TGZ:
                  tarPath = os.path.join(lockerDir, payload.name)
                  try:
                     with open(tarPath, 'wb') as newfobj:
                        shutil.copyfileobj(fobj, newfobj, length=payload.size)
                     Misc.extractTar(tarPath, pkgDir, exception=ValueError)
                  finally:
                     os.remove(tarPath)
               else:
                  msg = ('Locker payload %s of type %s is not supported'
                         % (payload.name, payload.payloadtype))
                  raise NotImplementedError(msg)
      db.Save()
