#!/usr/bin/python
########################################################################
# Copyright (c) 2010-2024 VMware, Inc.                                 #
# All Rights Reserved                                                  #
########################################################################

import os

from .. import IS_PATCH_THE_PATCHER, PYTHON_VER_STR
from ..ImageManager.Constants import INFO
from ..Utils.HostInfo import GetEsxVersionPair
from ..Version import VibVersion
from ..Vib import GetHostSoftwarePlatform

LIFECYCLE_DIR = os.path.join(os.path.sep, "var", "vmware", "lifecycle")

def _getRamdiskRoot():
   """Returns the root path for live, quickpatch, and boot installers to create
      ramdisks in.
   """
   # Legacy /tmp location.
   TMP_DIR = os.path.join(os.path.sep, 'tmp')
   # Current /usr/lib location.
   USR_LIB_DIR = os.path.join(os.path.sep, 'usr', 'lib', 'vmware', 'lifecycle')

   if not IS_PATCH_THE_PATCHER:
      # If not Patch the Patcher, no need to get the current version, use
      # USR_LIB_DIR directly.
      return USR_LIB_DIR

   version, patch = GetEsxVersionPair()
   # VSI only provides the main version and the patch level on 7.0.
   # Use a pseudo version string "version-patch" for comparison.
   curVer = VibVersion.fromstring(version + '-' + patch)
   esx803Ver = VibVersion.fromstring('8.0.3-0')

   if curVer < esx803Ver:
      # Upgrade from before the location change, use TMP_DIR.
      return TMP_DIR

   # Upgrade from after the location change.
   return USR_LIB_DIR

RAMDISK_ROOT = _getRamdiskRoot()


class Installer(object):
   """Installer is the base class behind all the *Installer classes
      and serves as an interface.  It does not implement any real
      installation functionality.

      Attributes:
         * database - A Database instance representing the package
                      database associated with the installer
   """
   SUPPORTED_VIBS = set()
   SUPPORTED_PAYLOADS = set()

   def __init__(self, task=None):
      """Constructor for this Installer class.
         Should determine if this class of installer should be created
         within the runtime environment.
         If so, initialize any databases.
         Parameters:
            * task - a vLCM task to update notifications/progress.
         Exceptions:
            InstallerNotAppropriate - Environment is not appropriate for
               this type of installer.
      """
      self._platform = GetHostSoftwarePlatform()
      self._task = task

   def __del__(self):
      """Destructor. Should always clean up after installations.
      """

   def StartTransaction(self, imgprofile, **kwargs):
      """Initiates a new installation transaction. Calculate what actions
         need to be taken.  Prepare installation destination as necessary.
         This method may change the installation destination.

         Parameters:
            * imgprofile - The ImageProfile instance representing the
                           target set of VIBs for the new image
            * 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
            * imgstate   - One of the HostImage.IMGSTATE constants
         Returns:
            A tuple (installs, removes), each of which is a 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 there is nothing to do, (None, None) is returned.
         Exceptions:
            InstallationError
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   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.

         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 - Set to True to get/put the payload from/to 
                           basemisc.tgz
         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.
            InstallationError
      """
      assert (read or write) and not (read and write), \
         "Only one of read or write can and must be True"

   def VerifyPayloadChecksum(self, vibid, payload):
      """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
      """

   def UpdateVibDatabase(self, newvib):
      """Update missing properties of vib metadata

         Parameters:
            * newvib   - The new vib to use as source
         Returns:
            None if the update succeeds, Exception otherwise
         Exceptions:
            VibFormatError
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def Cleanup(self, **kwargs):
      """Cleans up after a Transaction that has been started, perhaps from
         a previous instantiation of the Installer class.
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def CompleteStage(self):
      """Do what is needed to complete the stage operation.

         Exceptions:
            InstallationError
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def GetStagePath(self):
      """Get the path to the staged contents
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def Remediate(self):
      """Carry out the remediation operation.

         Returns:
            A Boolean, True if a reboot is needed.
         Exceptions:
            InstallationError
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def PostRemediationCheck(self):
      """Check after remediation.
         Override this method to perform checks in an installer.
      """
      pass

   @classmethod
   def GetSupportedQuickPatchPayloads(cls, vib, scriptsOnly=False):
      """Return a list of VIB payload names supported by the quick patch
         installer.

         Parameters:
            * vib - A VIB instance.
            * scriptsOnly - A boolean that indicates whether to return just the
                            lean payload containing the scripts or also include
                            patch payloads. Patch payloads are used during
                            remediation.
      """
      payloads = []
      scriptPayload = None
      for p in vib.payloads:
         if p.payloadtype in cls.SUPPORTED_PAYLOADS and \
            p.isQuickPatchRelevant(scriptsOnly=scriptsOnly):
            if p.isquickpatch:
               scriptPayload = p   # Only 1 script payload per vib
            else:
               payloads.append(p)
      payloads.sort(key=lambda x: x.overlayorder)
      payloads.insert(0, scriptPayload)

      return list(p.name for p in payloads)

   @classmethod
   def GetSupportedVibPayloads(cls, vib):
      """Return a set of VIB payload names supported by the installer. In case
         of quick patch, a sorted set based on overlayorder is returned.
         Parameters:
            * vib - A VIB instance.
      """
      return set(p.name for p in vib.payloads if p.payloadtype
                 in cls.SUPPORTED_PAYLOADS)

   def GetSupportedVibs(self, vibs):
      '''Return a set of VIB IDs which are supported by the installer.
         Parameters:
            * vibs - A VibCollection instance.
      '''
      return set(vib.id for vib in vibs.values() if vib.vibtype in
                 self.SUPPORTED_VIBS)

   def GetInstallationSize(self, imgprofile):
      '''Return total size of payloads by supported by this installer.
         Parameter:
            An imageprofile
         Returns:
            Total byte size of the payloads supported by this installer
      '''
      totalsize = 0
      for vibid in imgprofile.vibIDs:
         vib = imgprofile.vibs[vibid]

         # If the installer defines a platform, skip a non-applicable VIB.
         if self._platform is not None and not vib.HasPlatform(self._platform):
            continue

         if vib.vibtype in self.SUPPORTED_VIBS:
            for payload in vib.payloads:
               if payload.payloadtype in self.SUPPORTED_PAYLOADS:
                  totalsize += payload.size
      return totalsize

   def GetInstallerImageProfile(self, imgProfile):
      """Return the trimmed image profile for the installer based on the input:
         1) Unsupported VIBs are removed.
         2) Components that include unsupported VIBs are reserved.
         In practice, reserving unsupported components makes sure that even if
         the locker installer is offline or is not synced due to a rollback, the
         image profiles in live/bootbank are valid.
      """
      newProfile = imgProfile.Copy()
      rmVibIds = newProfile.vibIDs - self.GetSupportedVibs(newProfile.vibs)
      rmComps = newProfile.components.GetComponentsFromVibIds(rmVibIds)
      orphanRmVibIds = rmVibIds - rmComps.GetVibIDs()

      for comp in rmComps.IterComponents():
         # Takes care of both the component and VIBs in it.
         newProfile.ReserveComponent(comp.id)

      for vibId in orphanRmVibIds:
         newProfile.RemoveVib(vibId)

      return newProfile

   def GetImageProfileVibDiff(self, imgProfile):
      """Given new image profile to install, calculate installer-specific VIB
         diffs and return a tuple of sets: (adds, removes, keeps).
         This method assumes database/new image profiles both have removed
         unsupported types of VIBs (e.g locker).
      """
      # During an upgrade, current image profile may have partial reserved VIBs.
      curProfile = self.database.profile.GetPlatformSpecificCopy(
         self._platform, partialReservedVibs=True)

      newProfile = imgProfile.GetPlatformSpecificCopy(
         self._platform)
      adds, removes = newProfile.Diff(curProfile)
      return set(adds), set(removes), curProfile.vibIDs - set(removes)

   def SaveDatabase(self):
      """Write out the updated database of the installer.
      """
      raise NotImplementedError("Must instantiate a subclass of Installer")

   def UpdateImageProfile(self, newProfile):
      """Update image profile in the database when the VIB inventory managed by
         this installer has not changed but metadata has changed, i.e. base
         image, addon and components.
         This should be only used on live and bootbank installers as locker
         does not have an image profile.
      """
      oldProfile = self.database.profile

      imgProfile = self.GetInstallerImageProfile(newProfile)
      if not oldProfile.HasSameInventory(imgProfile) or \
         oldProfile.acceptancelevel != imgProfile.acceptancelevel:
         # There are metadata changes or the acceptance level has changed,
         # update the database with the image profile.
         self.database.PopulateWith(imgProfile=imgProfile)
         self.SaveDatabase()

   def AddTaskNotification(self, notificationId, msgArgs=None, resArgs=None,
                           type_=INFO):
      """Forms and adds a notification to the vLCM task.
      """
      if self._task is None:
         # No-op when no task is given.
         return

      # This will import VAPI modules not available on 6.7.
      from ..ImageManager.Utils import getNotification

      n = getNotification(notificationId, notificationId, msgArgs=msgArgs,
                          resArgs=resArgs, type_=type_)
      self._task.updateNotifications([n])
