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

"""Utility module for upgrade ESXi.

This module must maintain compatibility with two previous onprem releases
with regards to the ESXi environment and external dependencies.
"""

import json
import logging
import os
import random
import shutil
import time

import vmkctl
from coredump import getCoredumpPartition, unconfigureCoredumpToPart
from esxutils import EsxcliError, ramdiskAdd, ramdiskRemove
from vmSyslogUtils import ExecError, execCommand

from systemStorage import FS_TYPE_VFAT, FS_TYPE_VMFS, MiB, errors
from systemStorage.autoPartition import autoPartition
from systemStorage.bootbank import getBootbanksForInstall
from systemStorage.esxfs import fsRescan, getFssVolumes, umountFileSystems
from systemStorage.vmfsl import findSystemVmfsl
from traceback import format_exc

UPGRADE_BACKUP_PATH = '/tmp/upgBackups'

BOOTCFG_BACKUP_NAME = 'bootcfg.bak'

PARTEDUTIL_PATH = '/bin/partedUtil'
VMKFSTOOLS_PATH = '/bin/vmkfstools'
GPT_PART_TABLE = 'gpt'

GPT_VMK_DIAGNOSTIC_GUID = '9D27538040AD11DBBF97000C2911D1B8'
GPT_BASIC_DATA_GUID = 'EBD0A0A2B9E5443387C068B6B72699C7'
MBR_VMK_DIAGNOSTIC_PARTID = 252
MBR_VFAT_PARTID = 6

# From lib/vfat/vfat_int.h
VFAT_MAX_SECTORS = 128 * 65524

TEMP_BOOTBANK_SIZE = 500 * MiB

log = logging.getLogger(os.path.basename(__file__))

class TempBbSource(object):
   """Class to denote original source of tempBootBank
   """
   COREDUMP = "coredump"
   FREESPACE = "freespace"
   UNFORMATTED_VFAT = "unformattedVfat"
   SCRATCH = "scratch"

class TempBbMarker(dict):
   """Class to handle tempBootBank marker file operations.
   """
   # tempBootBank marker file
   MARKER_FILE = 'tempbbmarker.cfg'

   # tempBootBank marker file keys
   SOURCE = 'source'
   BOOTBANK = 'bootbank'
   ALTBOOTBANK = 'altbootbank'
   TEMPBOOTBANK = 'tempbootbank'
   PARTNUM = 'partNum'
   SCRATCH = 'scratch'

   def __init__(self, *args, **kwargs):
      self._validKeys= [TempBbMarker.TEMPBOOTBANK, TempBbMarker.SOURCE,
                   TempBbMarker.BOOTBANK, TempBbMarker.ALTBOOTBANK,
                   TempBbMarker.PARTNUM, TempBbMarker.SCRATCH]
      for key, value in kwargs.items():
         assert(key in self._validKeys)
      super().__init__(*args, **kwargs)

   def writeFile(self):
      """Writes tempBootBank metadata key values to marker file
      """
      filepath = os.path.join(self[TempBbMarker.TEMPBOOTBANK],
                              TempBbMarker.MARKER_FILE)
      try:
         with open(filepath, 'w') as f:
            json.dump(self, f)
      except Exception:
         log.exception('failed to write marker file: %s', filepath)
         raise

   def readFile(self, filepath):
      """Reads tempBootBank key values from the marker file
      """
      try:
         with open(filepath) as f:
            data = json.load(f)
         self.update(data)
      except Exception:
         log.exception("failed to read marker file: %s", filepath)
         raise

def removeTempBbMarker(tempBbMarker):
   """Delete the tempBootBank marker file in case of successful upgrade
   """
   tempBbVolumePath = tempBbMarker.get(TempBbMarker.TEMPBOOTBANK)
   tempBbMarkerFile = os.path.join(tempBbVolumePath, TempBbMarker.MARKER_FILE)
   if os.path.exists(tempBbMarkerFile):
      os.remove(tempBbMarkerFile)

def checkConfigUpgradeFailure():
   """Checks if there were configuration upgrade failures during esx upgrade
   """
   FAILED_UPGRADE_MODULE_FILE = '/var/run/vmware/files/failedUpgradeModules.txt'
   return os.path.isfile(FAILED_UPGRADE_MODULE_FILE)

def cleanFolder(dirPath, filesOnly=False):
   """Clean a folder.

   With filesOnly set, only remove files directly in the folder.
   """
   if not os.path.exists(dirPath):
      raise ValueError('Folder %s not found' % dirPath)
   for item in os.scandir(dirPath):
      path = os.path.join(dirPath, item.name)
      if item.is_dir():
         if not filesOnly:
            shutil.rmtree(path)
      else:
         os.remove(path)

def dirTreeToLower(dirRoot):
   """Convert any uppercase directories and files in the directory tree to
      lowercase.
   """
   renamePaths = []
   for root, dirs, files in os.walk(dirRoot):
      for d in dirs + files:
         if d.isupper():
            # Insert at head of array to rename deepest path first since
            # os.walk will not rescan the directory tree if a parent
            # directory gets renamed.
            renamePaths.insert(0, (os.path.join(root, d),
                                   os.path.join(root, d.lower())))

   for fromDir, toDir in renamePaths:
      try:
         # Rename to a temp path before the final one because some
         # filesystems, ie FAT, do not allow direct upper->lower conversion.
         tmp = '%s_%04d.tmp' % (fromDir, random.randrange(10000))
         os.rename(fromDir, tmp)
         os.rename(tmp, toDir)
      except Exception as e:
         log.warn('Error renaming %s to %s: %s', fromDir, toDir, str(e))

def copyFiles(srcPath, destPath, filesOnly=False, forceLowercaseName=False):
   """Copy files inside scrPath to destPath.

   With filesOnly set, only copy files and skip folders.

   @param forceLowercaseName: If True, destination filename will be converted
                              to lowercase if the source is all caps.
   """
   SVI_DIR_NAME = "System Volume Information"

   if not os.path.exists(srcPath):
      raise ValueError('Source folder %s not found' % srcPath)
   os.makedirs(destPath, exist_ok=True)

   for item in os.scandir(srcPath):
      src = os.path.join(srcPath, item.name)
      destFileName = item.name.lower() if (forceLowercaseName and
                                           item.name.isupper()) else item.name
      dest = os.path.join(destPath, destFileName)
      if item.is_dir():
         if not filesOnly and not item.name.startswith(SVI_DIR_NAME):
            # Ignore SVI directory created by Windows on DD, otherwise copytree
            # raises exception.
            try:
               shutil.copytree(src, dest,
                        ignore=shutil.ignore_patterns(SVI_DIR_NAME),
                        dirs_exist_ok=True)
            except TypeError:
               # On python < 3.8, dirs_exist_ok does not exist, retry
               # without the parameter.
               shutil.copytree(src, dest,
                        ignore=shutil.ignore_patterns(SVI_DIR_NAME))
      else:
         shutil.copy2(src, dest)

   if forceLowercaseName:
      dirTreeToLower(destPath)

def calculateDirMiBSize(path, filesOnly=False):
   """Calculate the total size of files inside a folder in MiB.

   With filesOnly set, only calculate files in the folder.
   """
   totalSize = 0
   if not filesOnly:
      for curPath, _, files in os.walk(path):
         for f in files:
            totalSize += os.path.getsize(os.path.join(curPath, f))
   else:
      for item in os.scandir(path):
         if item.is_file():
            totalSize += os.path.getsize(item.path)
   return totalSize // (1024 * 1024) + 4

def ramdiskRemoveWrapper(path, raiseException=True):
   """Wrap ramdiskRemove util to provide file deletion as well.

   When raiseException=False, ignore exceptions and try to perform both steps.
   """
   try:
      ramdiskRemove(path)
   except EsxcliError:
      if raiseException:
         raise
   try:
      shutil.rmtree(path)
   except EnvironmentError:
      if raiseException:
         raise

def getLockerAndScratchPaths(diskName=None):
   """Locate legacy SCRATCH and LOCKER partitions.

   Return the filesystem paths (/vmfs/volumes path) of the legacy locker and
   scratch partitions on the boot disk. Return None for a partition that is
   not found.
   """
   # Locker is 286 MiB, but we don't want to match exactly.
   LOCKER_MIN_SIZE = 280 * MiB
   LOCKER_MAX_SIZE = 290 * MiB
   # Scratch is 4GB, but there should be no other VFAT larger than 3GiB.
   SCRATCH_MIN_SIZE = 3072 * MiB

   if diskName is None:
      diskName, _ = getBootDiskInfo()

   lockerPath = scratchPath = None
   for vfatFsPtr in vmkctl.StorageInfoImpl().GetVFATFileSystems():
      vfatFs = vfatFsPtr.get()
      vfatFsHeadPart = vfatFs.GetHeadPartition().get()
      if vfatFsHeadPart.GetDeviceName() == diskName:
         fsPath = vfatFs.GetConsolePath()
         partSize = vfatFsHeadPart.GetSize()
         if partSize > LOCKER_MIN_SIZE and partSize < LOCKER_MAX_SIZE:
            lockerPath = fsPath
         elif partSize > SCRATCH_MIN_SIZE:
            scratchPath = fsPath

   log.debug('The boot disk contains locker %s and scratch %s.',
             lockerPath, scratchPath)
   return lockerPath, scratchPath

def _getBootVolume():
   """Get the bootbank volume for this boot.

   In case of loadEsx, we could have booted with a prefix in boot.cfg pointed to
   the temp bootbank used.
   """
   from vmware.esximage.Utils import BootCfg
   bootVol = vmkctl.SystemInfoImpl().GetBootVolume()
   bootCfgPath = os.path.join(bootVol, 'boot.cfg')
   bootCfg = BootCfg.BootCfg(bootCfgPath)
   if bootCfg.prefix != '' and os.path.exists(bootCfg.prefix):
      log.debug('Using prefix in %s as the boot volume: %s.',
                bootCfgPath, bootCfg.prefix)
      bootVol = bootCfg.prefix
   return bootVol

def upgradeBackup(bootBankPath=None, lockerPath=None, scratchPath=None):
   """Backup bootbank, locker, log (in scratch) and lifecycle dir (in scratch)
      of a disk to ramdisks.

   If no paths are given, the system boot device will be used.
   """
   AUDIT_LOG_FILE = 'audit.001'

   if bootBankPath is None:
      bootBankPath = _getBootVolume()
   if lockerPath is None:
      # Locker path is required, while scratch path can be None on
      # USB. Assume both need to be fetched when locker is not given.
      lockerPath, scratchPath = getLockerAndScratchPaths()

   logPath = os.path.join(scratchPath, 'log') if scratchPath and \
             os.path.isdir(os.path.join(scratchPath, 'log')) else None
   vmwarePath = os.path.join(scratchPath, 'vmware') if scratchPath and \
             os.path.isdir(os.path.join(scratchPath, 'vmware')) else None
   lifecyclePath = os.path.join(scratchPath, 'lifecycle') if scratchPath and \
             os.path.isdir(os.path.join(scratchPath, 'lifecycle')) else None

   # Add audit records in scratch for backup. These can be in any directory,
   # but always has the audit.001 file.
   auditPaths = []
   if scratchPath is not None:
      for (dirpath, dirs, files) in os.walk(scratchPath):
         if AUDIT_LOG_FILE in files:
            auditPaths.append(dirpath)

   # The temporary bootbank could be the scratch, thus count only the files
   # directly inside the partition.
   try:
      ramdiskSize = calculateDirMiBSize(bootBankPath, filesOnly=True) + \
                    calculateDirMiBSize(lockerPath) + 5
      if logPath is not None:
         ramdiskSize += calculateDirMiBSize(logPath)
      if vmwarePath is not None:
         ramdiskSize += calculateDirMiBSize(vmwarePath)
      if lifecyclePath is not None:
         ramdiskSize += calculateDirMiBSize(lifecyclePath)
      for p in auditPaths:
         ramdiskSize += calculateDirMiBSize(p, filesOnly=True)
   except Exception as e:
      raise errors.UpgradeBackupError('An error occurred while backing up VFAT '
                                      'partition files before re-partitioning: '
                                      'Failed to calculate size for temporary '
                                      'Ramdisk: %s. This might be due to VFAT '
                                      'corruption, please refer to KB 91136 '
                                      'for steps to remediate the issue.'
                                      % e, format_exc())

   log.info('Backing up bootbank %s, locker %s, log %s and lifecycle dir %s of '
            'the boot disk.', bootBankPath, lockerPath, logPath, lifecyclePath)

   try:
      ramdiskAdd(os.path.basename(UPGRADE_BACKUP_PATH),
                 UPGRADE_BACKUP_PATH, 0, ramdiskSize, '755')
   except Exception as e:
     raise errors.UpgradeBackupError('An error occurred while backing up VFAT '
                                     'partition files before re-partitioning: '
                                     'Failed to add Ramdisk of size %u: %s.'
                                     % (ramdiskSize, e), format_exc())

   try:
      copyFiles(bootBankPath, os.path.join(UPGRADE_BACKUP_PATH, 'bootbank'),
                filesOnly=True)
      copyFiles(lockerPath, os.path.join(UPGRADE_BACKUP_PATH, 'locker'),
                forceLowercaseName=True)
      if logPath is not None:
         copyFiles(logPath, os.path.join(UPGRADE_BACKUP_PATH, 'log'))
      if vmwarePath is not None:
         copyFiles(vmwarePath, os.path.join(UPGRADE_BACKUP_PATH, 'vmware'))
      if lifecyclePath is not None:
         copyFiles(lifecyclePath, os.path.join(UPGRADE_BACKUP_PATH,
                                               'lifecycle'))

      for p in auditPaths:
         # Retain directory hierarchy of audit directories
         subdir = p.replace(scratchPath + '/', '', 1)
         copyFiles(p, os.path.join(UPGRADE_BACKUP_PATH, 'audit', subdir),
                   filesOnly=True)
   except Exception as e:
      ramdiskRemoveWrapper(UPGRADE_BACKUP_PATH, raiseException=False)
      raise errors.UpgradeBackupError('An error occurred while backing up VFAT '
                                      'partition files before re-partitioning: '
                                      'Failed to copy files to Ramdisk: %s. '
                                      'This might be due to VFAT corruption, '
                                      'please refer to KB 91136 for steps to '
                                      'remediate the issue.' % e,
                                      format_exc())

def upgradeRestore(bootBankPath, osDataPath, lockerPath):
   """Restore backups of bootbank, locker, logs and /var/vmware contents.
   """
   copyFiles(os.path.join(UPGRADE_BACKUP_PATH, 'bootbank'), bootBankPath)
   copyFiles(os.path.join(UPGRADE_BACKUP_PATH, 'locker'), lockerPath)
   logBackupPath = os.path.join(UPGRADE_BACKUP_PATH, 'log')
   vmwareBackupPath = os.path.join(UPGRADE_BACKUP_PATH, 'vmware')
   lifecycleBackupPath = os.path.join(UPGRADE_BACKUP_PATH, 'lifecycle')
   auditBackupPath = os.path.join(UPGRADE_BACKUP_PATH, 'audit')
   if osDataPath is not None:
      if os.path.exists(logBackupPath):
         newLogPath = os.path.join(osDataPath, 'log')
         copyFiles(logBackupPath, newLogPath)
      if os.path.exists(vmwareBackupPath):
         newVmwPath = os.path.join(osDataPath, 'vmware')
         copyFiles(vmwareBackupPath, newVmwPath)
      if os.path.exists(lifecycleBackupPath):
         newLifecyclePath = os.path.join(osDataPath, 'vmware', 'lifecycle')
         copyFiles(lifecycleBackupPath, newLifecyclePath)
      if os.path.exists(auditBackupPath):
         copyFiles(auditBackupPath, osDataPath)

def removeBackups():
   """Remove backed up ramdisks.
   """
   ramdiskRemoveWrapper(UPGRADE_BACKUP_PATH)

def getBootDiskInfo():
   """Return the boot disk (name, full /vmfs path) tuple..

   Return (None, None) if the system is not booted from a disk.
   """
   bootDevice = vmkctl.SystemInfoImpl().GetBootDevice()
   if not bootDevice:
      log.debug('Boot disk not found, the system is likely PXE-booted.')
      return None, None
   bootDevicePath = '/vmfs/devices/disks/' + bootDevice
   if os.path.exists(bootDevicePath):
      log.debug('Boot disk is %s.', bootDevice)
      return bootDevice, bootDevicePath
   else:
      raise RuntimeError('Expected disk path %s does not exist'
                         % bootDevicePath)

def getLargeCoredumpNum(partTableType, parts, sectorSize):
   """Return the large goredump partition number.

   Check if disk includes a coredump (diagnostic) partition larger than the
   space required for a temp bootbank, which means the large coredump partition
   (9) if created by the installer/auto-partition. Returns the partition number.
   """
   partType = GPT_VMK_DIAGNOSTIC_GUID if partTableType == GPT_PART_TABLE \
              else str(MBR_VMK_DIAGNOSTIC_PARTID)

   for part in parts:
      if part[3] == partType:
         partSize = (int(part[2]) - int(part[1])) * sectorSize
         if partSize >= TEMP_BOOTBANK_SIZE:
            # Not exactly the size of the large curedump (2.5 GiB),
            # but our goal is to throw out the legacy small coredump
            # partition and stay capable to find any manually
            # created coredump partitions (unlikely).
            log.debug('Found large coredump partition %u.', int(part[0]))
            return int(part[0])
   return None

def getSystemScratchInfo():
   """Return scratch partition information.

   Info is returned as a (device name, filesystem path, partition number,
   free space) tuple, or (None, None, 0, 0) if /scratch is not configuted
   to a VFAT partition.
   """
   SCRATCH_LINK = '/scratch'
   scratchFsPath = os.readlink(SCRATCH_LINK)

   for vfatFsPtr in vmkctl.StorageInfoImpl().GetVFATFileSystems():
      vfatFs = vfatFsPtr.get()
      vfatHeadPart = vfatFs.GetHeadPartition().get()
      if vfatFs.GetConsolePath() == scratchFsPath:
         scratchDevice = vfatHeadPart.GetDiskLun().get().GetName()
         scratchFsPartNum = str(vfatHeadPart.GetPartition())
         freeSpace = vfatFs.GetSize() - (vfatFs.GetBlocksUsed() *
                                         vfatFs.GetBlockSize())
         log.debug('Scratch %s on disk partition %s:%s has %u bytes free.',
                   scratchFsPath, scratchDevice, scratchFsPartNum, freeSpace)
         return scratchDevice, scratchFsPath, scratchFsPartNum, freeSpace

   log.debug('The system does not have an active scratch partition.')
   return None, None, 0, 0

def isCoredumpConfigured(devicePath, partNum):
   """Check if the large coredump partition is configured to take coredumps.
   """
   partName = "%s:%u" % (os.path.basename(devicePath), partNum)
   return getCoredumpPartition() == partName

def getPartInfo(devicePath):
   """Return the partition table type (MBR/GPT) and partition entries.

   Each partition entries include the partition number, start sector, end
   sector, type/guid and attributes.
   """
   getCmd = '%s getptbl %s' % (PARTEDUTIL_PATH, devicePath)
   _, out, _ = execCommand(getCmd)

   # Parse partedUtil return and modify the requested partition,
   # each list of parts contains partNum, startSector, endSector,
   # GUID, attr.
   lines = out.split('\n')
   partTableType = lines[0].strip()
   parts = []
   for line in lines[2:]:
      part = line.strip().split(' ')
      if len(part) >= 5:
         parts.append(part)
   if partTableType == GPT_PART_TABLE:
      # For GPT remove the type string to keep output consistent
      for part in parts:
         del part[4]
   return (partTableType, parts)

def setPartTable(devicePath, partTableType, parts):
   """Set partition table using partedUtil, with 3 retries.

   Parts input is a list of list each containing: partition number, start
   sector, end sector, guid, attributes as required by partedUtil.
   """
   setCmd = '%s setptbl %s %s' % (PARTEDUTIL_PATH, devicePath, partTableType)
   for part in parts:
      setCmd += ' "%s" ' % ' '.join(part)

   for i in range(1, 4):
      try:
         execCommand(setCmd)
         break
      except ExecError as e:
         log.warn('Failed to set partition table on try %u: %s', i, str(e))

         try:
            # FS volumes that we need to use when logging opened files on a
            # failure. As this module should be useful for legacy layout only,
            # no need to consider VMFS-L.
            diskName = os.path.basename(devicePath)
            vfatVols = getFssVolumes(diskName=diskName, fsType=FS_TYPE_VFAT)
            vmfsVols = getFssVolumes(diskName=diskName, fsType=FS_TYPE_VMFS)

            for volume in vfatVols + vmfsVols:
               openFiles = volume.getOpenFiles()
               if openFiles:
                  log.debug("%s: volume has opened file handles: %s",
                            volume.name, str(openFiles))
         except Exception:
            pass

         if i == 3:
            raise
         time.sleep(3)

def countMbrPrimaryParts(parts):
   """Return the number of primary partition(s) in the MBR partition table.
   """
   EXT_PART_TYPE = 5
   count = 0
   lastExtEnd = 0 # Last end sector of an extended part.
   for part in parts:
      # Parts returned by partedUtil are sorted by layout (start/end sectors).
      partType = int(part[3])
      partStart = int(part[1])
      partEnd = int(part[2])
      if partType == EXT_PART_TYPE:
         # Extended partition, update last sector and count as
         # one primary part.
         lastExtEnd = partEnd
         count += 1
      elif partStart > lastExtEnd:
         # When a part starts after the last extended partition,
         # it is a primary part.
         count += 1
   return count

def formatVFATPartition(partPath):
   """Create a VFAT partition on a given partition.

   Return the volume path to the new partition.
   """
   import re
   cmd = '%s %s %s' % (VMKFSTOOLS_PATH, '-C vfat', partPath)
   _, out, _ = execCommand(cmd)

   p = re.compile('Successfully created new volume: ([A-Fa-f0-9-]{35})')
   m = p.search(out)
   if not m:
      raise RuntimeError('Failed to find the new VFAT volume name, stdout: %s'
                         % out)
   log.debug('VFAT volume %s has been formatted.', m.group(1))
   return '/vmfs/volumes/' + m.group(1)

def createNewVFATPart(diskPath, startSector, endSector, partTableType, parts):
   """Create a new VFAT partition at the end of the disk.

   On GPT, the partition number for the new VFAT partition is assigned by adding
   1 to the largest existing partition number.

   On MBR, we use any free primary partition number, assuming the disk has at
   most 3 primary partitions.

   Input:
      diskPath            - filesystem path to the disk
      startSector         - starting sector of the partition
      endSector           - ending sector of the partition
      partTableType/parts - current partition type and partition list
                            returned by getPartInfo()
   Return the /vmfs path and partition number to the new volume.
   """
   partNums = set(int(part[0]) for part in parts)

   if partTableType == GPT_PART_TABLE:
      partNum = max(partNums) + 1
      partId = GPT_BASIC_DATA_GUID
   else:
      freePrimaryNums = set(range(1, 5)) - partNums
      if freePrimaryNums:
         partNum = min(freePrimaryNums)
      else:
         raise RuntimeError('%s: failed to create VFAT partition, disk already '
                            'contains 4 primary partitions.' % diskPath)
      partId = str(MBR_VFAT_PARTID)

   newPart = [str(partNum), str(startSector), str(endSector), partId, '0']
   parts.append(newPart)

   log.info('Creating new VFAT partition %u.', partNum)

   # Set partition table and format the partition
   setPartTable(diskPath, partTableType, parts)
   return formatVFATPartition('%s:%u' % (diskPath, partNum)), partNum

def formatPartToVFAT(devicePath, partNum, partTableType, parts):
   """Modify partition table and format a partition to VFAT.

   Return the /vmfs path and partition number to the new volume.
      devicePath          - path to the device to work on
      partNum             - partition number to format
      partTableType/parts - current partition type and partition list
                            returned by getPartInfo()
   """
   partType = GPT_BASIC_DATA_GUID if (partTableType == GPT_PART_TABLE) \
              else str(MBR_VFAT_PARTID)

   # Modify the partition type in the partition table
   foundPart = False
   partTblChanged = False
   for part in parts:
      if int(part[0]) == partNum:
         foundPart = True
         if part[3] != partType:
            partTblChanged = True
            part[3] = partType

         # Check partition length against VFAT max.
         startSector = int(part[1])
         endSector = int(part[2])
         if endSector - startSector + 1 > VFAT_MAX_SECTORS:
            # Resize the partition to end earlier.
            endSector = startSector + VFAT_MAX_SECTORS - 1
            log.info('Resize partition %u to end at sector %u', partNum,
                     endSector)
            partTblChanged = True
            part[2] = str(endSector)
         break

   if not foundPart:
      raise ValueError('Partition %u does not exist on device %s'
                       % (partNum, devicePath))

   log.info('Formatting existing partition %u to VFAT.', partNum)

   # Only write the partition table when modified.
   if partTblChanged:
      setPartTable(devicePath, partTableType, parts)

   # Format the volume
   return formatVFATPartition('%s:%u' % (devicePath, partNum)), partNum

def getUnformattedVFATPart(partTableType, parts, sectorSize):
   """Find a VFAT partition suitable for holding the temporary bootbank.

   Find any unformatted VFAT partition with enough space for the temporary
   bootbank. Return its partition number, or None if not partition is found.
   """
   partType = GPT_BASIC_DATA_GUID if partTableType == GPT_PART_TABLE \
              else str(MBR_VFAT_PARTID)

   candidates = set()
   for part in parts:
      if part[3] == partType:
         partSize = (int(part[2]) - int(part[1])) * sectorSize
         if partSize >= TEMP_BOOTBANK_SIZE:
            candidates.add(int(part[0]))

   # Unformatted partition will not appear in the file system list
   if candidates:
      for vfatFsPtr in vmkctl.StorageInfoImpl().GetVFATFileSystems():
         vfatPartNum = vfatFsPtr.get().GetHeadPartition().get().GetPartition()
         if vfatPartNum in candidates:
            candidates.remove(vfatPartNum)

   if candidates:
      # Only return one
      partNum = list(candidates)[0]
      log.info('Found unformatted VFAT partition %u.', partNum)
      return partNum
   else:
      return None

def getUnusedVFATVolume(diskName, sectorSize, requiredSize):
   """Find any unused VFAT volume of the specified size.

   For the boot disk, this means the volume is not symlinked as a bootbank,
   locker, or scratch. Likely, the volume was created as the temporary bootbank.
   Returns the vmfs volume path to the file system.
   """
   SYS_SYMLINKS = ['/bootbank', '/altbootbank', '/store', '/scratch']
   excludedPaths = [os.path.realpath(symlink) for symlink in SYS_SYMLINKS]

   # At worst case, we are 63 sector off the desired size due to alignment
   # of 64 sectors.
   graceSize = 63 * sectorSize

   for vfatFsPtr in vmkctl.StorageInfoImpl().GetVFATFileSystems():
      vfatFs = vfatFsPtr.get()
      vfatFsHeadPart = vfatFs.GetHeadPartition().get()
      fsPath = vfatFs.GetConsolePath()
      if vfatFsHeadPart.GetDeviceName() == diskName and \
         vfatFsHeadPart.GetSize() + graceSize >= requiredSize and \
         not fsPath in excludedPaths:
         log.info('Found existing available VFAT volume %s.', fsPath)
         return fsPath
   return None

def getTempBootbank(bootbank, altbootbank):
   """Find the temporary bootbank and create the marker file

   Find the temporary bootbank vfat partition on the boot disk
   and write related metadata to tempBootbank marker file which
   is used later for the purpose of rolling back to legacy 6.7
   layout in case of one or more config upgrade module failures.
   The tempBootBank marker is cleared otherwise.
   """
   source, tempBootbank, partNum = findTempBootbank()
   if source and tempBootbank and partNum:
      scratch = os.path.realpath('/scratch')
      bootbank = os.path.realpath(bootbank)
      altbootbank = os.path.realpath(altbootbank)
      params = {
         TempBbMarker.SOURCE: source,
         TempBbMarker.TEMPBOOTBANK: tempBootbank,
         TempBbMarker.BOOTBANK: bootbank,
         TempBbMarker.ALTBOOTBANK: altbootbank,
         TempBbMarker.PARTNUM: partNum,
         TempBbMarker.SCRATCH: scratch
      }
      try:
         tempBbMarker = TempBbMarker(**params)
         tempBbMarker.writeFile()
      except Exception:
         log.error("failed to create tempBootBank marker file")
   return tempBootbank

def findTempBootbank():
   """Find and format (when necessary) a temporary larger bootbank
      (a VFAT partition) on the boot disk.

   Returns:
      Source of the partition taken up by temporary bootbank, path to
      the VFAT filesystem prepared as the temporary bootbank and the
      corresponding partition number. The partition will always be emptied,
      except for scratch, only files in the partition are removed to keep
      folders, incl. log intact.

   The following methods are tried in order:
      1 If the disk already has a temporary bootbank created by method 2/3/4,
        simply clean the partition and return it for use.
      2 If the disk has more than 500 MiB in free space, create a VFAT partition
        use the free space. For MBR, we will not be able to simply append a
        partition when there are already 4 primary partitions on the disk.
      3 If a VFAT partition is listed in the partition table but yet not
        formatted. This is possible on partition tables created by late 6.0 and
        6.5 installers where a bug prevents the scratch partition from being
        formatted.
      4 If the disk has a coredump partition that is larger than 500 MiB (the
        larger coredump, partition 9, is typically 2.5GiB), deactivate coredump
        and format the partition to VFAT. For an auto-partitioned coredump, it
        will be resized if it exceeds the max FAT16 size.
      5 If the disk has a scratch partition (partition 2) and have more than
        500MiB free, use the partition directly.
   """
   diskName, diskPath = getBootDiskInfo()
   partTableType, parts = getPartInfo(diskPath)
   diskLun = vmkctl.StorageInfoImpl().GetDeviceByName(diskName).get()
   sectorSize = diskLun.GetBlockSize()

   # Check if we have already created a temporary bootbank by the first
   # three strategies.
   ununsedVfatVol = getUnusedVFATVolume(diskName, sectorSize,
                                        TEMP_BOOTBANK_SIZE)
   if ununsedVfatVol is not None:
      cleanFolder(ununsedVfatVol)
      return None, ununsedVfatVol, None

   # Now proceed to create/find a temporary bootbank.
   # First, try to find free space.
   lastSector = diskLun.GetNumBlocks() - 1
   lastPartEndSector = int(parts[-1][2])
   freeBytes = (lastSector - lastPartEndSector) * sectorSize
   if partTableType == GPT_PART_TABLE:
      # Exclude secondary GPT, one sector for the header, and then 128
      # partition entries each with 128 bytes.
      GPT_ARRAY_SIZE = 128 * 128
      freeBytes -= GPT_ARRAY_SIZE + sectorSize
      # For a system disk, we assume we are well below 128 partitions
      # and thus can add a new partition as long as there is enough
      # space.
      canCreatePart = True
   else:
      # MBR can append a new primary partition when there are less
      # than 4 existing primary ones.
      canCreatePart = countMbrPrimaryParts(parts) <= 3
   if freeBytes >= TEMP_BOOTBANK_SIZE and canCreatePart:
      endSector = lastPartEndSector + TEMP_BOOTBANK_SIZE // sectorSize
      volPath, partNum = createNewVFATPart(diskPath, lastPartEndSector + 1,
                                           endSector, partTableType, parts)
      return TempBbSource.FREESPACE, volPath, partNum

   # Second, check for a VFAT partition that is unformatted.
   rawVfatPart = getUnformattedVFATPart(partTableType, parts, sectorSize)
   if rawVfatPart is not None:
      volPath, partNum = formatPartToVFAT(diskPath, rawVfatPart,
                                          partTableType, parts)
      return TempBbSource.UNFORMATTED_VFAT, volPath, partNum

   # Third, check for a larger coredump.
   largeCoredumpPart = getLargeCoredumpNum(partTableType, parts, sectorSize)
   if largeCoredumpPart is not None:
      if isCoredumpConfigured(diskPath, largeCoredumpPart):
         # Disable coredump to partition when the partition is configured.
         log.info('Disabling active coredump partition %u.', largeCoredumpPart)
         unconfigureCoredumpToPart()
      volPath, partNum = formatPartToVFAT(diskPath, largeCoredumpPart,
                                          partTableType, parts)
      return TempBbSource.COREDUMP, volPath, partNum

   # Finally, if scratch is configured on the boot disk and have enough free
   # space, use it as the temporary bootbank.
   scDeviceName, scFsPath, scFsPartNum, scFreeSpace = getSystemScratchInfo()
   if scDeviceName == diskName and scFreeSpace >= TEMP_BOOTBANK_SIZE:
      # Remove free-floating files in the scratch before using it to store
      # bootbank modules. Log and other folders are kept.
      cleanFolder(scFsPath, filesOnly=True)
      return TempBbSource.SCRATCH, scFsPath, scFsPartNum

   raise RuntimeError('Failed to find an eligible temporary bootbank')

def repartitionFromLegacy(disk, bootbankPath, lockerPath, scratchPath):
   """Upgrade the ESX partition table from 6.x to 7.x layout.

   XXX: Rather than letting Weasel re-partition the ESX boot disk before reboot,
        we should only have it stage the new 7.x system in a temporary bootbank
        like esximage does for live upgrades, and then rely on jumpstart to
        perform AutoPartition post-reboot.

        In addition to unifying the code paths between ISO-install and live
        upgrade, that would also help to support roll-back.
   """
   # backup to RAM
   log.info("%s: backing up boot disk to RAM (bootbank=%s, locker=%s, "
            "scratch=%s)", disk.name, bootbankPath, lockerPath, scratchPath)
   try:
      upgradeBackup(bootbankPath, lockerPath, scratchPath)
   except errors.UpgradeBackupError as e:
      log.error("%s\n%s" % (e.msg, e.stackTrace))
      raise Exception(e.msg)

   # re-partition
   log.info("%s: umounting VFAT/VMFS partitions", disk.name)
   umountFileSystems(disk.name)

   log.info("%s: re-partitioning boot disk (keepDatastore=True)", disk.name)
   autoPartition(disk, keepDatastore=True, createDatastore=False)
   fsRescan()

   # Restore from RAM
   bootbankPath = getBootbanksForInstall(disk.name)[0].path
   vmfslPath = findSystemVmfsl(disk.name).path
   log.info("%s: restoring boot disk contents from RAM (bootbank=%s, "
            "VMFSL=%s)", disk.name, bootbankPath, vmfslPath)
   osdataPath = vmfslPath if disk.supportsOsdata else None
   lockerPath = (os.path.join(osdataPath, "locker") if osdataPath is not None
                                                    else vmfslPath)
   upgradeRestore(bootbankPath, osdataPath, lockerPath)

def backupBootCfg(bootbankPath):
   """Backup boot.cfg file from the specified bootbank path.
   """
   src = os.path.join(bootbankPath, 'boot.cfg')
   dst = os.path.join(bootbankPath, BOOTCFG_BACKUP_NAME)
   if os.path.isfile(src):
      shutil.copyfile(src, dst)

def restoreBootCfg(bootbankPath):
   """Restore /boot.cfg from the backup in the bootbank path if it exists.
   """
   src = os.path.join(bootbankPath, BOOTCFG_BACKUP_NAME)
   dst = os.path.join(bootbankPath, 'boot.cfg')
   if os.path.isfile(src):
      shutil.copyfile(src, dst)
