#!/usr/bin/python

'''
Copyright 2018-2020 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 os
import shutil
import tarfile
from tempfile import TemporaryDirectory

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

from systemStorage.ddImage import DDImage
import systemStorage.vfat as vfat


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

   def __init__(self, imageProfile, mmdBin, mcopyBin, vfatSo):
      """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
      vfat.libvfat = vfat.loadVfatLibrary(vfatSo)

   def Write(self, imgFilename, checkdigests=True, checkacceptance=True,
             kernelopts=None, isARM=False):
      """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
            * isARM: Set to true to write the ARM image.
      """
      with TemporaryDirectory() as stageDir:

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

         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)

         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'))
         self._AddReservedVibs(os.path.join(hostSeedDir, 'reservedVibs'))

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

   def _GetBootbankFiles(self, bootbank1Dir, bootbank2Dir, kernelopts,
                         checkdigests, checkacceptance):
      """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
      """
      self._GetBootbankVibs(bootbank1Dir, checkdigests, checkacceptance)
      self._CreateBootCfg(bootbank1Dir, True, kernelopts)
      self._CreateDatabase(bootbank1Dir)

      self._CreateBootCfg(bootbank2Dir, False, kernelopts)

   def _GetBootloaderFiles(self, bootloaderDir, checkdigests, checkacceptance,
                           isARM):
      """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
            * isARM: Set to true to write the ARM image.
      """
      self._CheckVibFiles(checkacceptance=checkacceptance)

      for vibid in self.imageprofile.vibIDs:
         vibObj = self.imageprofile.vibs[vibid]
         if vibObj.vibtype != Vib.BaseVib.TYPE_BOOTBANK:
            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 = {'BOOT%s.EFI' % ('AA64' if isARM else 'x64'): '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):
      """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

      """
      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()

      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 vibObj.vibtype == Vib.BaseVib.TYPE_BOOTBANK:
            for payload, fobj in vibObj.IterPayloads(checkDigests=checkdigests):
               if payload.payloadtype in payload.NON_GZIP_TYPES:
                  # 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 _CreateBootCfg(self, bootCfgDir, isPrimaryBootbank, kernelopts):
      """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
      """
      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)

      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)
                     with tarfile.open(tarPath) as tar:
                        tar.extractall(pkgDir)
                  finally:
                     os.remove(tarPath)
               else:
                  msg = ('Locker payload %s of type %s is not supported'
                         % (payload.name, payload.payloadtype))
                  raise NotImplementedError(msg)
      db.Save()
