#!/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 logging
import os
import tempfile
import tarfile

from .. import Depot, Downloader, Errors, Vib
from ..ImageManager import ReservedVibTar
from ..Utils import BootCfg, EsxGzip
from ..Utils.HashedStream import HashedStream
from ..Utils.Misc import isString, seekable

"""This module provides an ancestor class for creating different image types
   based on an image profile."""

log = logging.getLogger("ImageBuilder")

def resetFObj(fObj):
   """Resets a file object or an HashedStream object.
   """
   if isinstance(fObj, HashedStream):
      fObj.reset()
   else:
      fObj.seek(0)

def getSeekableFObj(fObj, tmpDir=None):
   """Returns a seekable file object based on the given one, a temporary
      file will be created if it is not seekable.
   """
   PAYLOAD_READ_CHUNKSIZE = 1024 * 1024

   if seekable(fObj):
      return fObj

   tfp = tempfile.TemporaryFile(dir=tmpDir)
   bytesRead = fObj.read(PAYLOAD_READ_CHUNKSIZE)
   while bytesRead:
      tfp.write(bytesRead)
      bytesRead = fObj.read(PAYLOAD_READ_CHUNKSIZE)
   fObj.close()
   tfp.seek(0)

   return tfp

def createTgz(payload, payloadName, tarDest, mode=None):
   """Creates a tgz containing a single file.
         Parameters:
            * payload     - File object of the payload to be tgz'ed.
            * payloadName - Name to be given to the payload in the tgz.
            * tarDest     - Complete path with file name of the tgz to
                            be created.
            * mode        - Additional file mode(s) to set.
   """
   if isString(tarDest):
      fobj = EsxGzip.GzipFile(tarDest, "wb")
   else:
      fobj = EsxGzip.GzipFile(fileobj=tarDest, mode="wb")

   with tarfile.open(fileobj=fobj, mode="w",
                     format=tarfile.GNU_FORMAT) as depotTar:
      tarinfo = tarfile.TarInfo(payloadName)

      payload = getSeekableFObj(payload)
      payload.seek(0, 2)
      tarinfo.size = payload.tell()
      if mode:
         tarinfo.mode = tarinfo.mode | mode
      resetFObj(payload)

      depotTar.addfile(tarinfo, payload)

   fobj.close()

class ImageBuilder(object):
   """This class is a skeleton to be inherited for the different methods of
   generating an image with the contents of an image profile."""

   DATABASE_NAME = "imgdb.tgz"
   PAYLOADTAR_NAME = "imgpayld.tgz"
   # Payloads of esx-base that are not directly stored in bootbanks. They are
   # tar'ed in this file and mounted as one.
   BASE_MISC_PAYLOADTAR_NAME = 'basemisc.tgz'
   RESERVED_VIBS_TAR_NAME = 'resvibs.tgz'
   ESXIO_DEPOT_TAR_NAME = 'esxiodpt.tgz'
   ESXIO_KS_TAR_NAME = 'esxio_ks.tgz'
   ESXIO_ENABLED = os.getenv('ESX_IO', 'False') == 'True'

   def __init__(self, imageprofile, vibdownloadfn=Depot.VibDownloader):
      """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.
               * vibdownloadfn - A callable with the signature
                                `(destfile: str, vibobj: Vib.Vib) -> None`
                                which, given a vib object `vibobj`, will write
                                the vib into the given destination path
                                `destfile`.
      """
      self.imageprofile = imageprofile
      self.vibdownloadfn = vibdownloadfn

   @staticmethod
   def _GetSeekableFileObj(url):
      # Each VIB can potentially contain more than one payload. It is likely
      # that the order of payloads within the VIB will not be the order in
      # which they will be written to the target. Therefore, the file-like
      # object from which a VIB's payloads are read must be seekable. This
      # method takes a URL or path as input and returns a seek-able file object.
      # For remote URLs, it accomplishes that by downloading to a temporary
      # file.
      if os.path.exists(url):
         # Avoid using Downloader and all the python/ssl dependencies for
         # a local path, especially during build.
         d = open(url, 'rb')
      else:
         from .. import Downloader
         d = Downloader.Downloader(url).Open()
      # We only care if something is seekable.
      if hasattr(d, "seek"):
         return d
      # Else, copy it locally.
      t = tempfile.TemporaryFile()
      data = d.read(512)
      while data:
         t.write(data)
         data = d.read(512)
      d.close()
      t.seek(0)
      return t

   @classmethod
   def _CheckVibFile(cls, vib, checkacceptance=True):
      # This method just ensures that a single VIB has an appropriate file-like
      # object behind it.

      # Should probably make public attributes/properties/methods to implement
      # this check.
      if vib._arfile is not None and vib._arfile.seekable():
         pass
      elif not vib.remotelocations:
         msg = ("VIB %s has neither a seek-able file object nor a URL "
                "location. This may indicate a problem with the depot "
                "metadata. " % vib.id)
         raise Errors.VibIOError(None, msg)
      else:
         problems = list()
         success = False
         for url in vib.remotelocations:
            try:
               vib.OpenFile(cls._GetSeekableFileObj(url))
               # We only need one. :-)
               success = True
               break
            except Exception as e:
               problems.append(str(e))

         if not success:
            msg = ("Error retrieving file for VIB '%s': %s."
                   % (vib.id, "; ".join(problems)))
            raise Errors.VibDownloadError(None, None, msg)

      if checkacceptance:
         vib.VerifyAcceptanceLevel()

   def _CheckVibFiles(self, checkacceptance=True, platform=None,
                      partialDepot=False):
      # This method makes sure that every VIB has an appropriate file object
      # behind it, so that we'll be able to read payloads into the target object.
      # It also checks to make sure that things in vibIDs have corresponding
      # objects in vibs. This is the only place where we make this check, so
      # this method should be called before _AddPayloads() or _AddDatabase().
      for vibid in self.imageprofile.vibIDs:
         try:
            vib = self.imageprofile.vibs[vibid]
         except KeyError:
            msg = "Could not find object for VIB '%s'." % vibid
            raise Errors.ProfileFormatError(self.imageprofile.name, msg)
         # This fills in VIB signature, required even when VIB is not for the
         # platform.
         try:
            self._CheckVibFile(vib, checkacceptance)
         except Errors.VibDownloadError:
            # When the source depot is partial, VIBs of other platform are not
            # present. Ignore only those cases.
            if partialDepot and not vib.HasPlatform(platform):
               pass
            else:
               raise

      # We need to make sure that every reserved VIB has an appropriate
      # file object behind it.
      for vibid in self.imageprofile.reservedVibIDs:
         try:
            vib = self.imageprofile.reservedVibs[vibid]
         except KeyError:
            msg = "Could not find object for reserved VIB '%s'." % vibid
            raise Errors.ProfileFormatError(self.imageprofile.name, msg)
         if not platform or vib.HasPlatform(platform):
            self._CheckVibFile(vib, checkacceptance)

   def _GetBootCfg(self, installer=True, moduleroot='', isoImage=False,
                   kernelopts=None, bootbankVibOnly=False,
                   appendResVibsTgz=True, esxiodepot=None, esxioKsfile=None,
                   platform=None, extraBootMods=None):
      '''Return BootCfg instance if boot modules is not zero, otherwise return
         None
         Parameters:
            * installer  - True if the bootcfg is for installer
            * moduleroot - root for module files
            * isoImage   - True if the bootcfg is used to build an iso
            * kernelopts - Additional kernel boot options other than
                           feature states
            * bootbankVibOnly - if True, only include bootbank VIB modules.
            * appendResVibsTgz - if True, resvibs.tgz is appended to modules.
            * esxiodepot - If not None, ESXIO_DEPOT_TAR_NAME is appended to
                           modules.
            * platform   - SoftwarePlatform productLineID whose VIB payloads
                           are to be written to the ISO. VIBs for other
                           platforms are ignored but metadata is still
                           included.
            * extraBootMods - Optional list of additional boot modules
      '''
      payload_types = [Vib.Payload.TYPE_TGZ, Vib.Payload.TYPE_VGZ,
                       Vib.Payload.TYPE_BOOT, Vib.Payload.TYPE_VZSTD]
      if installer:
         payload_types.append(Vib.Payload.TYPE_INSTALLER_VGZ)

      vib_types = [Vib.BaseVib.TYPE_BOOTBANK] if bootbankVibOnly else None
      bootorder = self.imageprofile.GetBootOrder(payload_types, vib_types,
                     platform=platform)

      modules = [p.localname for (vibid, p) in bootorder]
      if not modules:
         return None

      bootcfg = BootCfg.BootCfg()

      # ESX Database
      modules.append(self.DATABASE_NAME)
      # Misc esx-base payloads.
      modules.append(self.BASE_MISC_PAYLOADTAR_NAME)
      if extraBootMods:
         modules += [os.path.basename(path) for path in extraBootMods]
      if appendResVibsTgz:
         modules.append(self.RESERVED_VIBS_TAR_NAME)
      if esxiodepot:
         modules.append(self.ESXIO_DEPOT_TAR_NAME)
      if esxioKsfile:
         modules.append(self.ESXIO_KS_TAR_NAME)
      if installer:
         bootcfg.kernelopt["runweasel"] = None
         modules.append(self.PAYLOADTAR_NAME)
         bootcfg.title = "Loading ESXi installer"
      else:
         bootcfg.kernelopt["autoPartition"] = "TRUE"
         bootcfg.title = "Loading ESXi"

      if isoImage:
         bootcfg.kernelopt["cdromBoot"] = None

      if kernelopts is not None:
         bootcfg.kernelopt.update(kernelopts)

      if moduleroot:
         modules = [os.path.join(moduleroot, module) for module in modules]

      if self.ESXIO_ENABLED:
         bootcfg.runtimewd = 1
         bootcfg.runtimewdtimeout = 300

      bootcfg.kernel = modules[0]
      bootcfg.modules = modules[1:]
      bootcfg.build = self.imageprofile.GetEsxVersion()

      return bootcfg

   def _AddReservedVibs(self, reservedVibTarPath, platform=None,
                        createEmptyTar=True):
      """This method generates a tar file that contains reserved vibs,
         filter by platform if specified.

         Parameters:
            * reservedVibTarPath - A string containing the path where the tar
                                   file will be created or a file-like object
                                   in which the tar will be produced.
            * createEmptyTar     - If True, only add an empty resvibs.tgz and
                                   skip all reserved VIBs. Defaults to False.
      """
      reservedVibTar = ReservedVibTar.ReservedVibTar(reservedVibTarPath)
      try:
         if createEmptyTar:
            log.info('Creating empty reserve VIB tar - resvibs.tgz')
            return

         for vibid in self.imageprofile.reservedVibIDs:
            # Skip adding VIB in resvibs.tgz if it is currently installed
            if vibid in self.imageprofile.vibIDs:
               log.info('Skip adding vib %s in resvibs.tgz as it is '
                        'currently installed.', vibid)
               continue

            vib = self.imageprofile.reservedVibs[vibid]

            # skip esxio vibs in esxi and vice versa
            if platform and not vib.HasPlatform(platform):
               continue

            # Windows does not allow double opening of a file, close the
            # initial descriptor and cleanup after use.
            with tempfile.NamedTemporaryFile(delete=False) as tmpfd:
               localPath = tmpfd.name
            try:
               self.vibdownloadfn(localPath, vib)
               relativevibpath = vib.relativepath or vib.GetRelativePath()
               reservedVibTar.AddVib(localPath, relativevibpath)
            finally:
               os.remove(localPath)
      except Exception as e:
         msg = "Could not download and package reserved VIBs. %s" % str(e)
         raise Errors.VibDownloadError(None, None, msg)
      finally:
         reservedVibTar.close()

   def _AddPayloads(self, target, checkdigests=True):
      # This method adds each payload to <target>, while doing the right things
      # for various payload types.
      raise NotImplementedError("_AddPayloads is not implemented in the child"
                                " class.")

   def _AddDatabase(self, target):
      # This method generates a tar database from the image profile, writes it
      # to a temp file, then adds the temp file to <target>.
      raise NotImplementedError("_AddDatabase is not implemented in the child"
                                " class.")

   def _AddMetadataZip(self, target):
      # This method generates a metadata.zip from the image profile.
      raise NotImplementedError("_AddMetadataZip is not implemented in the"
                                " child class.")

   def _AddProfileXml(self, target):
      # Adds profile.xml to the upgrade directory in the target.
      raise NotImplementedError("_AddProfileXml is not implemented in the"
                                " child class.")

   # Not necessary for a pxeboot, but could be generated either way just incase
   # it is needed.
   def _AddBootCfg(self, target, installer=True):
      # This method populates efi/boot/boot.cfg. It is required for EFI.
      raise NotImplementedError("_AddBootCfg is not implemented in the child"
                                " class")

