#!/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
from enum import IntEnum
from .. import Errors, Vib, ZSTD_COMPRESSION_ENABLED
from ..Utils import HostInfo, PathUtils
from ..Utils.Misc import LogLargeBuffer
from . import Installer
from . import InstallerCommon as Ic

log = logging.getLogger('LiveImageInstaller')

class FileState(IntEnum):
   """Flags for type of file operation.
   """
   keepreg = 1
   keepoverlay = 2
   removereg = 4
   removeoverlay = 8
   addreg = 16
   addoverlay = 32

class LiveImageInstaller(Installer):
   '''LiveImageInstaller is the Installer class to live install/remove VIBs for
      live system.

      Attributes:
         * database - A Database.Database instance of the live system
   '''

   installertype = "live"
   priority = 5

   # Support locker VIB for PXE host so live install/remove of the tools
   # VIB is supported and the VIB will be kept in the image profile.
   if HostInfo.IsPxeBooting():
      SUPPORTED_VIBS = set([Vib.BaseVib.TYPE_BOOTBANK,
                            Vib.BaseVib.TYPE_LOCKER,])
   else:
      SUPPORTED_VIBS = set([Vib.BaseVib.TYPE_BOOTBANK,])

   SUPPORTED_PAYLOADS = set([Vib.Payload.TYPE_TGZ, Vib.Payload.TYPE_VGZ])
   # When the FSS ESXZstdCompression is enabled, SUPPORTED_PAYLOADS will contain
   # TYPE_VZSTD.
   if ZSTD_COMPRESSION_ENABLED:
      SUPPORTED_PAYLOADS.add(Vib.Payload.TYPE_VZSTD)

   BUFFER_SIZE = 8 * 1024

   def __init__(self, root='/', quickPatch=False, task=None, **kwargs):
      super().__init__(task=task)
      self.liveimage = Ic.LiveImage(root, quickPatch=quickPatch)
      self.problems = list()

   def GetStagePath(self):
      if self.liveimage:
         return self.liveimage.STAGE_DIR
      else:
         return None

   def isImgProfileStaged(self, imgprofile):
      """Returns a boolean to indicate whether the given image profile is
         staged. If there is some other profile in the stage database, it will
         be cleaned up before returning False.
         Parameters:
            * imgprofile - The new image profile that is being staged or
                           remediated.
      """
      if self.stagedatabase is not None:
         if self.stagedatabase.profile and \
               self.stagedatabase.profile.vibIDs == imgprofile.vibIDs:
            return True
         self.liveimage.Cleanup()
      return False

   @property
   def isstaged(self):
      return self.liveimage.isstaged

   @property
   def database(self):
      return self.liveimage.database

   @property
   def stagedatabase(self):
      return self.liveimage.stagedatabase

   @property
   def stagedimageprofile(self):
      """
      Returns the staged image profile if present
      """
      stagedatabase = self.stagedatabase
      if stagedatabase:
         return stagedatabase.profile
      return None

   def VerifyPrerequisites(self, imgstate=None, forcebootbank=False):
      """Verifies if the options passed and the current state of system
         allows rebootless installation.
      """

      from ..HostImage import HostImage
      if forcebootbank:
         msg = 'Nothing to do for live install - live installation has been ' \
               'disabled.'
         self.problems.append(msg)
         log.debug(msg)

      if imgstate not in (HostImage.IMGSTATE_FRESH_BOOT,
            HostImage.IMGSTATE_LIVE_UPDATED):
         msg = ('Only reboot-required installations are possible right now as '
                'a reboot-required installation has been done.')
         self.problems.append(msg)
         log.debug(msg)

      if self.database.profile is None:
         msg = 'No ImageProfile found for live image.'
         raise Errors.InstallationError(None, None, msg)

   def VerifyTransaction(self, imgprofile, adds, removes, keeps):
      """Verify whether there are any issues that prevent live installation.
         Returns True is there are problems, False otherwise.
      """
      problems = self._CheckTransaction(imgprofile, adds, removes, keeps)
      if problems:
         log.debug('The transaction is not supported for live install:\n')
         LogLargeBuffer(str(problems), log.debug)
         self.problems = problems
         return True
      return False

   def StartTransaction(self, imgprofile, imgstate=None, preparedest=True,
                        forcebootbank=False, stageonly=False, **kwargs):
      """Initiates a new installation transaction. Calculate what actions
         need to be taken. Note that quick patch related works will only be
         executed if called from subclass QuickPatchInstaller.

         This method only works on staging directory

         Parameters:
            * imgprofile  - The ImageProfile instance representing the
                            target set of VIBs for the new image
            * imgstate    - The state of current HostImage, one of IMGSTATE_*
            * preparedest - Boolean, if True, then prepare the destination.
                            Set to false for a "dry run", to avoid changing
                            the destination.
            * forcebootbank - Boolean, if True, skip install of live image
                              even if its eligible for live install
            * stageonly - Boolean, if True, only stages the contents, i.e.
                          changes are made only to stageliveimage
         Returns:
            A StartTransactionResult instance that has attributes: installs,
            removes, staged.
            Installs and removes are list of VIB IDs for HostImage.Stage()
            to install to the destination and to remove from the destination,
            in order to make it compliant with imgprofile.

            If LiveImage has already staged the imgprofile, staged is True.

            If there is nothing to do, StartTransactionResult(None, None, False)
            is returned.
         Exceptions:
            InstallationError
      """
      unsupportedRes = Ic.StartTransactionResult(None, None, False)

      # Skip if reboot required VIBs have been installed

      self.VerifyPrerequisites(imgstate=imgstate, forcebootbank=forcebootbank)
      if self.problems:
         return unsupportedRes

      imgprofile = self.GetInstallerImageProfile(imgprofile)
      adds, removes, keeps = self.GetImageProfileVibDiff(imgprofile)

      # Return unsupported in live installer if quickpatch installer has staged
      # an image that cannot be regularly live installed.
      haveProblems = self.VerifyTransaction(imgprofile, adds, removes, keeps)
      if haveProblems:
         return unsupportedRes

      staged = self.isImgProfileStaged(imgprofile)
      res = Ic.StartTransactionResult(adds, removes, staged)

      if staged:
         return res

      imgsize = self.GetInstallationSize(imgprofile)
      if preparedest and (removes or adds):
         self.liveimage.StartTransaction(imgprofile, imgsize,
                                         stageonly=stageonly)

      return res

   def VerifyPayloadChecksum(self, vibid, payload, **kwargs):
      """Verify the checksum of a given payload.

         Parameters:
            * vibid   - The Vib id containing the payload
            * payload - The Vib.Payload instance to read or write
         Returns:
            None if verification succeeds, Exception otherwise
         Exceptions:
            ChecksumVerificationError
            InstallationError
      """
      if payload.payloadtype not in self.SUPPORTED_PAYLOADS:
         # TODO: we are currently only checking the SUPPORTED_PAYLOADS
         # for this installer. Once we have the ability to lookup
         # modules hashes from the kernel, we should enable checking
         # for other boot modules.
         return

      try:
         if vibid not in self.database.profile.vibstates:
            msg = ('Could not locate VIB %s in LiveImageInstaller'
                   % vibid)
            raise Errors.InstallationError(None, [vibid], msg)
         vibstate = self.database.profile.vibstates[vibid]

         if payload.name not in vibstate.payloads:
            msg = "Payload name '%s' of VIB %s not in LiveImage DB" % \
                  (payload.name, vibid)
            raise Errors.InstallationError(None, [vibid], msg)
         tardiskname = vibstate.payloads[payload.name]
         Ic.VerifyLiveTardiskChecksum(tardiskname, payload)
      except Exception as e:
         msg = "Failed to verify checksum for payload {}: {}" \
            .format(payload.name, e)
         log.error("%s", msg)
         raise

   def UpdateVibDatabase(self, newvib):
      """Update missing properties of vib metadata
         New vibs are always installed in the liveimage

         Parameters:
            * newvib   - The new vib to use as source
         Returns:
            None if the update succeeds, Exception otherwise
         Exceptions:
            VibFormatError
      """
      self.liveimage._UpdateVib(newvib)

   def OpenPayloadFile(self, vibid, payload, read=True, write=False,
                       isBaseMisc=False):
      """Creates and returns a File-like object for either reading from
         or writing to a given payload.  One of read or write must be True, but
         read and write cannot both be true.

         Parameters:
            * vibid   - The Vib id containing the payload
            * payload - The Vib.Payload instance to read or write
            * read    - Set to True to get a File object for reading
                        from the payload.
            * write   - Set to True to get a File object for writing
                        to the payload.
            * isBaseMisc - Not used, defaults to False.
         Returns:
            A File-like object, must support read (for read), write (for
            write), close methods.
            None if the desired read/write is not supported.
         Exceptions:
            AssertionError    - neither read nor write is True, or both are true
            InstallationError - Cannot open file to write or read
      """
      if isBaseMisc:
         raise NotImplementedError("Do not know how to handle isBaseMisc")

      Installer.OpenPayloadFile(self, vibid, payload, read, write)

      if read:
         if vibid not in self.database.profile.vibstates:
            msg = 'Could not locate VIB %s in LiveImageInstaller' % vibid
            raise Errors.InstallationError(None, [vibid], msg)

         vibstate = self.database.profile.vibstates[vibid]
         if payload.name not in vibstate.payloads:
            msg = "Payload name '%s' of VIB %s not in LiveImage DB" % (
                  payload.name, vibid)
            raise Errors.InstallationError(None, [vibid], msg)

         if (payload.payloadtype in self.SUPPORTED_PAYLOADS or
             payload.payloadtype == Vib.Payload.TYPE_BOOT):

            filepath = os.path.join(self.liveimage.tardisksdir,
                  vibstate.payloads[payload.name])
            if os.path.isfile(filepath):
               # Tardisk is not compressed
               return open(filepath, 'rb')
            else:
               msg = "Payload '%s' of VIB %s at '%s' is missing" % (
                     payload.name, vibid, filepath)
               if HostInfo.HostOSIsSimulator():
                  log.info("HostSimulator: %s", msg)
                  return None
               else:
                  raise Errors.InstallationError(None, [vibid], msg)
         else:
            return None
      else:
         if payload.payloadtype not in self.SUPPORTED_PAYLOADS:
            log.debug("Payload %s of type '%s' in VIB '%s' is not supported by "
                      "LiveImageInstaller.", payload.name, payload.payloadtype,
                      vibid)
            return None

         filepath = os.path.join(self.liveimage.stagedatadir, payload.localname)
         try:
            # Decompressed file stream is required when writing to this stream
            return open(filepath, 'wb')
         except EnvironmentError as e:
            msg = 'Can not open %s to write payload %s: %s' % (filepath,
                  payload.name, e)
            raise Errors.InstallationError(e, None, msg)

   def Cleanup(self, checkStaged=False):
      """Cleans up the live image staging area.
         The cleanup will only happen if:
            - checkStaged is False, regardless of whether the image is staged;
            - checkStaged is True and the image is staged.
         Parameter:
            * checkStaged - A boolean that indicates whether we need to check
                            the image is staged before doing the cleanup.
      """
      if checkStaged and not self.isstaged:
         return

      try:
         self.liveimage.Cleanup()
      except Exception:
         pass

   def CompleteStage(self):
      """Complete the staging of live image by writing out the database and
         staged indicator to staging directory.

         Exceptions:
            InstallationError
      """
      self.liveimage.CompleteStage()

   def Remediate(self, checkmaintmode=True, hasConfigDowngrade=False, **kwargs):
      """Live remove and install VIB payloads

         For each VIB to remove, shutdown scripts will be executed and payloads
         will be unmounted. For each VIB to install, payloads will be mounted
         and init scripts will be executed.

         Returns:
            A Boolean, always False, as a reboot is not needed.
         Exceptions:
            InstallationError -
            HostNotChanged    - If there is no staged ImageProfile
      """
      if not self.isstaged:
         msg = 'LiveImage is not yet staged, nothing to remediate.'
         raise Errors.HostNotChanged(msg)
      if self.problems:
         msg = 'Live installation is not supported or has been disabled, ' \
               'skip remediation.'
         raise Errors.HostNotChanged(msg)

      adds, removes, _ = \
         self.GetImageProfileVibDiff(self.liveimage.stagedatabase.profile)

      self.liveimage.Remediate(adds, removes, checkmaintmode,
                               hasConfigDowngrade, LiveImageInstaller)
      return Ic.RemediationResult(False, None)

   def _detectOverlayProblems(self, groups, unsupported, allowedOverlays=None):
      """Iterate the groups and classify the filepaths based on unsupported
         map to collect overlay problems.

         Parameters
            * groups: A tuple of (vibId, vibList, FileState flags)
            * unsupported: A map of key as FileState flags and value as error
            * allowedOverlays: A pair of a set of file paths and a tuple of
                               allowed overlay scenarios.
      """
      exceptionFiles, exceptionFlags = allowedOverlays if allowedOverlays else \
         (set(), ())

      overlaid = {
         FileState.keepoverlay | FileState.addreg : \
            ("File to be installed is overlaid by existing VIB", []),
         FileState.keepreg | FileState.addoverlay : \
            ("File to be installed overlays existing VIB", []),
         FileState.addreg | FileState.addoverlay : \
            ("File to be installed overlays another VIB to be installed", []),
      }
      problems = []

      # Build a map from file path to the file states. File state is a bit map
      # to indicate which group the file belongs to.
      filestates = {}
      for vibids, vibs, flags in groups:
         for vibid in vibids:
            vib = vibs[vibid]
            for filepath in vib.filelist:
               if filepath == '' or filepath.endswith('/'):
                     continue

               filepath = PathUtils.CustomNormPath('/' + filepath)
               ind = vib.overlay and 1 or 0
               if filepath in filestates:
                  filestates[filepath] |= flags[ind]
               else:
                  filestates[filepath] = flags[ind]

      for filepath in filestates:
         if filestates[filepath] in unsupported:
            normPath = os.path.normpath(filepath)
            if (normPath in exceptionFiles and
                filestates[filepath] in exceptionFlags):
               # The file is in the exception set and the overlay scenario is
               # allowed.
               continue
            problem = "%s : %s" % (unsupported[filestates[filepath]], filepath)
            problems.append(problem)
         if filestates[filepath] in overlaid:
            overlaid[filestates[filepath]][1].append(filepath)

      # Files overwritten are logged. This info is applicable no matter
      # unsupported or not.
      for overlay in overlaid:
         if overlaid[overlay][1]:
            log.info("%s : %s", overlaid[overlay][0],
                     overlaid[overlay][1])

      return problems

   def _checkLiveInstallRemove(self, imageprofile, adds, removes):
      """Check whether the VIBs to add or remove have liveinstallok or
         liveremoveok, respectively.
         Parameters:
            * imageprofile - The ImageProfile instance for the new image.
            * adds/removes - A set of VIB IDs of VIBs to add/remove.
         Returns:
            * problems - A list of problem strings, each of which represents
                         an issue, i.e., a VIB cannot be live installed or
                         removed.
      """
      problems = []
      for vibid in adds:
         # liveinstallok must be True for:
         # new VIB if called by LiveImageInstaller;
         # new non-QuickPatch VIB if called by QuickPatchInstaller.
         if not imageprofile.vibs[vibid].liveinstallok:
            problem = 'VIB %s cannot be live installed.' % vibid
            problems.append(problem)

      for vibid in removes:
         # liveremoveok must be True for VIB to be removed. If called by
         # QuickPatchInstaller, VIBs to remove have already deducted those
         # replaced by Quick Patch VIBs to install.
         if not self.database.vibs[vibid].liveremoveok:
            problem = 'VIB %s cannot be removed live.' % vibid
            problems.append(problem)

      return problems

   def _CheckTransaction(self, imageprofile, adds, removes, keeps):
      """Check the transaction to see if there are any logical reasons that
         prevent live installation of the transaction.
            * For VIBs to be removed: require liveremoveok
            * For VIBs to be installed: require liveinstallok
            * No VIBs to be installed are overlaid by existing VIB
            * No VIBs to be removed are overlaid by existing VIB
         Parameter:
            * imageprofile - The ImageProfile instance for the new image.
            * adds/removes/keeps - A set of VIB IDs of VIBs to add/remove/keep.
         Returns:
            * problems - A list of problem strings, each of which represents
                         a failure, i.e., a VIB cannot be live installed removed
                         or there is an unsupported file operation.
      """
      # Check overlay.
      groups = ((keeps, self.database.vibs, (FileState.keepreg,
                                             FileState.keepoverlay)),
                (removes, self.database.vibs, (FileState.removereg,
                                               FileState.removeoverlay)),
                (adds, imageprofile.vibs, (FileState.addreg,
                                           FileState.addoverlay)))

      unsupported = {
         FileState.keepoverlay | FileState.removereg  : \
            "File to be removed is overlaid by existing VIB",
         FileState.keepoverlay | FileState.removereg | FileState.addreg : \
            "File to be removed/installed is overlaid by existing VIB",
         FileState.keepoverlay | FileState.addreg : \
            "File to be installed is overlaid by existing VIB",
         #XXX: visorFS to support resurrect?
         FileState.keepreg | FileState.removeoverlay : \
            "File to be removed overlays existing VIB",
         FileState.keepreg | FileState.removeoverlay | FileState.addoverlay : \
            "File to be removed/installed overlays existing VIB",
      }

      problems = self._detectOverlayProblems(groups, unsupported)

      # Check live install or remove.
      problems += self._checkLiveInstallRemove(imageprofile, adds, removes)

      return problems

   def SaveDatabase(self):
      """Save live image database.
         Called in the base Installer class for image profile updates,
         regular transactions are handled by LiveImage.Remediate().
      """
      self.liveimage._UpdateDB(useStageDB=False)
