# Copyright 2018-2019,2022 VMware, Inc.
# All rights reserved. -- VMware Confidential

"""Functions relative to ESX bootbanks.
"""
import os

from systemStorage import BOOTBANK_LINK, ALTBOOTBANK_LINK, FS_TYPE_VFAT
from systemStorage.esxfs import EsxFsUuid, FssVolume, getFssVolumes
from vmware.esximage.Utils.BootCfg import BootCfg, BootCfgError
from vmware import vsi

MiB = 1024 * 1024

class EsxVersionInfo(object):
   """ESXi version information class for a bootbank.
   """
   def __init__(self, versionStr=None):
      self._versionStr = None
      self._majorVersionTuple = None
      self._isLatest = False
      if versionStr is not None:
         self.setEsxVersion(versionStr)

   @property
   def isLatest(self):
      """True if this bootbank has the highest 'updated' value of all bootbanks.
      """
      return self._isLatest

   @property
   def versionStr(self):
      """ESX version string.
      """
      return self._versionStr

   @property
   def majorVersionTuple(self):
      """Tuple of the major version digits, e.g. (7, 0, 0).
      """
      return self._majorVersionTuple

   def setEsxVersion(self, versionStr):
      """Verify and set ESX with a version string.
      """
      try:
         esxMajorVer = versionStr.split('-')[0]
         verTuple = tuple(map(int, esxMajorVer.split('.')))
      except ValueError:
         raise ValueError('Invalid ESXi version %s' % versionStr)
      self._versionStr = versionStr
      self._majorVersionTuple = verTuple

   def setLatest(self, isLatest):
      """Set isLatest property.
      """
      assert isinstance(isLatest, bool)
      self._isLatest = isLatest

class EsxBootbank(FssVolume):
   """A read-only class that contains bootbank information.
   """

   def __init__(self, vsiInfo):
      super().__init__(vsiInfo)

      self._bootCfg = None

      bootCfgPath = os.path.join(self.path, 'boot.cfg')
      if os.path.isfile(bootCfgPath):
         try:
            self._bootCfg = BootCfg(bootCfgPath)
         except BootCfgError:
            pass

   @property
   def bootCfg(self):
      """The bootCfg object represents the boot.cfg contents.

      Return None if /boot.cfg does not exist, or if the config file contains a
      'bootstate' value that is not success/updated.
      """
      return self._bootCfg

   @property
   def isValid(self):
      """True if this bootbank is in a valid bootstate.
      """
      if self._bootCfg is not None:
         return self._bootCfg.bootstate in (BootCfg.BOOTSTATE_SUCCESS,
                                            BootCfg.BOOTSTATE_UPDATED)
      return False

   @property
   def isDirty(self):
      """True if this bootbank is in dirty/attempted bootstate.
      """
      if self._bootCfg is not None:
         return self._bootCfg.bootstate == BootCfg.BOOTSTATE_ATTEMPTED
      return False

   @property
   def updatedSerial(self):
      """The bootbank's 'updated' value, or None if the bootbank is invalid.
      """
      if self.isValid:
         return self._bootCfg.updated
      return None

   @property
   def build(self):
      """ESX version string (e.g. '7.0.0-0.0.12345').
      """
      if self.isValid:
         return self._bootCfg.build
      return None

   @property
   def kernelOptStr(self):
      """The bootbank's 'kernelopts' value, or None if the bootbank is invalid.
      """
      if self.isValid:
         return self._bootCfg.kerneloptToStr(self._bootCfg.kernelopt)

   @property
   def isActive(self):
      """True if this is the bootbank ESX was booted from.
      """
      try:
         return self.uuid == getBootFsUUID()
      except Exception:
         return False

def getBootFsUUID():
   """Retrieve the boot filesystem UUID.
   """
   uuidStr = vsi.get("/system/bootFsUUID")["bootUuidStr"]
   return EsxFsUuid.fromString(uuidStr)

def getActiveBootbank():
   """Get the EsxBootbank object representing the current bootbank.

   The current bootbank is the one which the system booted from, i.e. the value
   of the kernel's bootFsUUID, which might not necessarily be symlinked with
   /bootbank.

   The returned value is an EsxBootbank object with the boot.cfg file
   parsed (if available).
   """
   try:
      # bootFsUUID hex byte array needs to be converted to UUID type
      uuid = getBootFsUUID()
   except Exception:
      return None

   volumes = getFssVolumes(uuid=uuid)
   if len(volumes) == 0:
      # bootUUID was provided in the boot option, but the FS does not exist
      return None
   assert len(volumes) == 1, "%u volumes with UUID %s" % (len(volumes), uuid)
   volume = volumes[0]

   return EsxBootbank.fromFssVolume(volume)

def identifyBootbanks(includeAll=False):
   """Return the FSS volumes pointed to by the /{alt}bootbank symlinks.

   If includeAll is True, then also include any other VFAT volume which contain
   /boot.cfg file.
   """
   if not includeAll and not os.path.exists(BOOTBANK_LINK):
      raise FileNotFoundError("Bootbank not configured on system")

   bbRealPath = os.path.realpath(BOOTBANK_LINK)
   altbbRealPath = os.path.realpath(ALTBOOTBANK_LINK)
   bankList = []

   for volume in getFssVolumes(fsTypes=[FS_TYPE_VFAT]):
      if os.path.exists(os.path.join(volume.path, 'boot.cfg')):
         if includeAll or volume.path in (bbRealPath, altbbRealPath):
            bank = EsxBootbank.fromFssVolume(volume)
            bankList.append(bank)

   return bankList

def getLatestBootbank():
   """Return the bootbank with the highest 'updated' value.

   An example use case is to identify which volume was last upgraded and will
   become active at next reboot.
   """
   nextUpdated = -1
   nextBootbank = None

   for bank in identifyBootbanks():
      updated = bank.updatedSerial
      if updated is not None and updated > nextUpdated:
         nextUpdated = updated
         nextBootbank = bank

   if nextBootbank is None:
      raise FileNotFoundError("No valid bootbank found")

   return nextBootbank

def scanEsxOnDisks(diskNames):
   """Find ESX boot volumes.

   Given a list of disk names, build a dictionary that contains ESXi version
   information. The dictionary is indexed by names of disks, and then by file
   system UUIDs.
   """
   esxVolumes = dict()

   for diskName in diskNames:
      latestUpdated = -1
      latestVol = None

      for volume in getFssVolumes(diskName=diskName, fsTypes=[FS_TYPE_VFAT]):
         bank = EsxBootbank.fromFssVolume(volume)
         updated = bank.updatedSerial
         if updated is not None:
            try:
               esxVersion = EsxVersionInfo(bank.build)
            except ValueError:
               # Ignore bootbanks with invalid bank.build string
               continue

            esxVolumes.setdefault(diskName, {})[volume] = esxVersion

            if updated > latestUpdated:
               latestUpdated = updated
               latestVol = volume

      if latestVol is not None:
         esxVolumes[diskName][latestVol].setLatest(True)

   return esxVolumes

def findBootbanks(diskName):
   """Find all bootbanks on a disk.

   Return a list of EsxBootbank objects for all bootbanks found on the given
   disk (regardless of whether they contatn a /boot.cfg). The list is sorted
   by partition numbers.

   This function works on both legacy (250MB bootbanks) and 7.0+ disk layouts,
   by looking at the partition layout, size and sectors.

   Prereq: filesystems on the disk must have been mounted prior to calling this
   function.
   """
   from systemStorage.esxdisk import EsxDisk

   OLD_BB_SIZE_MiB = 250
   OLD_BB_MIN_SIZE_MiB = 240

   bootbankList = []

   # Identify VMware FATs by scanning the partition's second sector for
   # the VMware magic signature. All FAT partitions created by ESXi and tools
   # have the magic string.
   disk = EsxDisk(diskName)
   disk.scanPartitions()
   vfatPartNums = disk.getVfatPartitions()

   # Legacy layout will have 4 or more VMware FAT partitions, with at least:
   # system boot, 2 bootbanks, and locker.
   # The legacy boot partition is a FAT16 partition (for 512-byte sectors)
   # with a different GPT/MBR partition type; hence ESXi does not mount it.
   isLegacy = len(vfatPartNums) > 3

   for volume in getFssVolumes(diskName=disk.name, fsTypes=[FS_TYPE_VFAT]):
      if volume.partNum not in vfatPartNums:
         # No point inspecting if not a VMware partition
         continue

      if (isLegacy and
          not (OLD_BB_MIN_SIZE_MiB <= volume.sizeInMB <= OLD_BB_SIZE_MiB)):
            continue

      bootbank = EsxBootbank.fromFssVolume(volume)
      bootbankList.append(bootbank)

   assert len(bootbankList) == 2, "len(bootbankList) == %u" % len(bootbankList)

   bootbankList.sort(key=lambda b: b.partNum)
   return bootbankList

def getBootbanksForInstall(diskName):
   """Return the FSS volumes for the 2 bootbanks on the given disk.

   Scan the disk to locate the current BOOTBANK and ALTBOOTBANK partitions, and
   return both partitions in reverse order (ALTBOOTBANK, BOOTBANK). if bootbanks
   are found but none is currently active, lookup the partitions labels and make
   'BOOTBANK1' the new active bootbank.

   Only bootbanks with a /boot.cfg config file and a valid 'bootstate' value are
   listed, otherwise None is returned.
   """
   oldBootbank = oldAltBootbank = None

   latestUpdated = 0
   bootbankList = findBootbanks(diskName)

   # Loop the two bootbanks to designate bootbank and altbootbank.
   for bootbank in bootbankList:
      if (bootbank.isValid and
          bootbank.updatedSerial > latestUpdated):
         if oldBootbank is not None:
            # This is the old bootbank, and we have seen the actual
            # altbootbank already, thus push previously found
            # bootbank to altbootbank.
            oldAltBootbank = oldBootbank
         oldBootbank = bootbank
         latestUpdated = bootbank.updatedSerial
         continue

      # No boot.cfg or invalid bootstate.
      if oldAltBootbank is None:
         # Update altbootbank but not the updated serial number.
         oldAltBootbank = bootbank
      else:
         # Both bootbanks are invalid (the previous bootbank has been
         # set as altbootbank). This happens when we fresh install,
         # or in a rare situation, with a damaged system disk.
         # As findBootbanks() has returned bootbanks sorted by part
         # numbers, we want to use the first one as the new bootbank,
         # thus it can stay as the "old altbootbank".
         oldBootbank = bootbank

   assert oldBootbank is not None and oldAltBootbank is not None
   # Swap them to get the next bootbanks.
   return oldAltBootbank, oldBootbank
