# Copyright (c) 2018-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.

"""Contain utilities for auto-partitioning ESX disks.

Depending on the availibility and type of storage devices, ESX System Storage is
partitioned as follow:

   +----+----+----------------------------+-----------+    \
   | BB | BB |           OSDATA           | DATASTORE |     |
  HD----+----+----------------------------+-----------+     |
                                                            |
   +----+----+--------+    +--------------+-----------+     | Persistent
   | BB | BB | unused |    |    OSDATA    | DATASTORE |     |   OSDATA
  USB---+----+--------+   HD--------------+-----------+     |
                                                            |
   +----+                  +--------------+-----------+     |
   | BB |                  |    OSDATA    | DATASTORE |     |
  RAM---+                 HD--------------+-----------+    /


   +----+----+--------+    +--------------+                \
   | BB | BB | LOCKER |    |    OSDATA    |                 |
  USB---+----+--------+   RAM-------------+                 | Non-persistent
                                                            |     OSDATA
   +----+                  +--------------+                 |
   | BB |                  |    OSDATA    |                 |
  RAM---+                 RAM-------------+                /

  Classes of storage devices:
    HD:   High-endurance, persistent
    USB:  Low-endurance, persistent
    RAM:  High-endurance, non-persistent

OSDATA vs. LOCKER

  LOCKER is a last resort partition which is only present if no persistent
  OSDATA can be created. LOCKER is guaranteed to be persistent, however it is
  low-endurance and limited in size. The partition is restricted to only 2 very
  specific usages: storing a local copy of the vmtools, and optionally
  collecting coredumps.
"""
import sys
from systemStorage import *
from systemStorage.vfat import *

if IS_ESX:
   from systemStorage.datastore import getNextDatastoreLabel
   from systemStorage.bootbank import getBootFsUUID
   from esxutils import getVmkBootOptions, sysAlert
   from systemStorage.vmfsl import getVmfsLabel, vmfsUnmount

OSDATA_MAX_END_MB = 128 * 1024                # 128GB

MEDIA_SIZES = {
   "min": 32 * 1024,
   "small": 64 * 1024,
   "default": OSDATA_MAX_END_MB,
   "max": sys.maxsize,
}

# Alignment by 1MB helpers for partition allocation
roundupMBSector = lambda x, sectorsPerMB: (x + sectorsPerMB) & ~(sectorsPerMB-1)
calcEndSector = lambda start, mb, sectorsPerMB: start + mb * sectorsPerMB - 1

def getMediaSize(sysMediaSize=None):
   """Get applicable system-storage media size option.

   If sysMediaSize is None, then return the default size.
   """
   if sysMediaSize is not None:
      sysMediaSize = sysMediaSize.lower()
      if sysMediaSize not in MEDIA_SIZES:
         raise ValueError("Invalid media size: %s" % sysMediaSize)
      return MEDIA_SIZES[sysMediaSize]

   try:
      opt = getVmkBootOptions()[SYSTEM_MEDIA_SIZE_OPT].lower()
      if opt in MEDIA_SIZES:
         return MEDIA_SIZES[opt]
   except Exception:
      pass
   return OSDATA_MAX_END_MB

def computeBootbankSize(disk, keepDatastore):
   """Calculate the optimal bootbank size (in MiB) for a given disk.

   The bootbank size varies with the amount of disk space available before any
   datastore partition. If no datastore is present (or if keepDatastore=False),
   the entire disk is considered in the calculation.
   """
   # PR 3355379: during esxcli upgrade, mounted systemStorage library does not
   # always properly import from the appropriate directory. So, have this as a
   # as a local import. In case of importError, set upgradeBootbankFSSEnabled = False.
   try:
      from systemStorage.upgradeUtils import checkUpgradeBootbankFSSEnabled
      upgradeBootbankFSSEnabled = checkUpgradeBootbankFSSEnabled()
   except ImportError:
      upgradeBootbankFSSEnabled = False
   freeSpace = disk.sizeInMB

   if keepDatastore:
      datastore = disk.getDatastore()
      if datastore is not None:
         freeSpace = datastore.start * disk.sectorSize // MiB

   if freeSpace >= SYSTEM_DISK_SIZE_LARGE_MB:
      return BOOTBANK_SIZE_LARGE_MB
   elif (freeSpace >= SYSTEM_DISK_SIZE_MEDIUM_MB and
         not upgradeBootbankFSSEnabled):
      return BOOTBANK_SIZE_MEDIUM_MB
   elif freeSpace >= SYSTEM_DISK_SIZE_SMALL_MB:
      return (BOOTBANK_SIZE_MEDIUM_MB if
              upgradeBootbankFSSEnabled else
              BOOTBANK_SIZE_SMALL_MB)
   else:
      raise ValueError("%s: not enough available disk space to create ESX boot "
                       "partitions (available=%uMiB < required=%uMiB)" %
                       (disk.name, freeSpace, SYSTEM_DISK_SIZE_SMALL_MB))

def addVfatPartitions(disk, bootbankSize, sectorsPerMB, bootPartSize=None,
                      bootbank1Size=None, bootbank2Size=None):
   """Set the VFAT partitions for the given disk but do not sync.

   @param bootbankSize: Size of use for the bootbanks. The value is overriden by
                        bootbank1Size and bootbank2Size if these are not None.
   @param bootPartSize: If not set, use default system boot partition size
   @param bootbank1Size: If not set, use value of bootbankSize
   @param bootbank2Size: If not set, use value of bootbankSize
   @return: boot, bootbank1 and bootbank2 partition objects.
   """
   bootPartSize = bootPartSize or SYSTEM_BOOT_SIZE_MB
   bootbank1Size = bootbank1Size or bootbankSize
   bootbank2Size = bootbank2Size or bootbankSize

   disk.clearPartitions()
   firstSector = 64 * 512 // disk.sectorSize
   end = calcEndSector(firstSector, bootPartSize, sectorsPerMB)

   boot = disk.setPartition(1, BOOTPART_FS, firstSector, end, BOOTPART_LABEL,
                            bootable=True)

   # Give each partition a gap between each other
   start = roundupMBSector(end + sectorsPerMB, sectorsPerMB)
   end = calcEndSector(start, bootbank1Size, sectorsPerMB)

   # Reuse the provided boot UUID from an upgrade and loadESX. On regular
   # rebooted upgrades and clean install, the UUID doesn't exist.
   try:
      bootFsUUID = str(getBootFsUUID())
   except Exception:
      bootFsUUID = None

   bootbank1 = disk.setPartition(5, BOOTBANK_FS, start, end, BOOTBANK1_LABEL,
                                 uuid=bootFsUUID)

   start = roundupMBSector(end + sectorsPerMB, sectorsPerMB)
   end = calcEndSector(start, bootbank2Size, sectorsPerMB)
   bootbank2 = disk.setPartition(6, BOOTBANK_FS, start, end, BOOTBANK2_LABEL)

   return boot, bootbank1, bootbank2

def autoPartition(disk, keepDatastore=False, createDatastore=True,
                  sysMediaSize=None):
   """Auto-partition the ESX boot disk.

   Parameters:
      keepDatastore    Whether to preserve an existing VMFS datastore;
                       has no effect if there is no datastore on the disk.
      createDatastore  Whether we want to create a datastore on the disk;
                       even when set to True, for a small disk or a USB
                       disk, there still will be no datastore created.
      sysMediaSize     System-storage media size option (one of MEDIA_SIZES key)
                       to limit the (supported) size of OSData.

   CAUTION: this method will wipe out your entire disk!
   Exception: with keepDatastore, VMFS datastore would be preserved.
   """
   sectorsPerMB = 1024 * 1024 // disk.sectorSize

   if keepDatastore:
      if disk.hasDatastore:
         # Get the datastore part before clearing the part table.
         datastore = disk.getDatastore()
      else:
         # Ignore @keepDatastore option if there is no datastore.
         keepDatastore = False

   bootbankSize = computeBootbankSize(disk, keepDatastore)
   _, _, bootbank2 = addVfatPartitions(disk, bootbankSize, sectorsPerMB)
   end = bootbank2.end
   start = roundupMBSector(end + sectorsPerMB, sectorsPerMB)

   if keepDatastore:
      vmfsStart = datastore.start
      vmfsEnd = datastore.end
      volumePath = "%s:%u" % (disk.path, datastore.num)
      try:
         vmfsLabel = getVmfsLabel(volumePath)
      except Exception:
         # VMFS partition is not yet formatted
         vmfsLabel = ""

      # Back up 2 MBs and align to the closest end of MB before the
      # datastore, this guarantees at least 1 MB of unoccupied space.
      osDataEnd = (roundupMBSector(vmfsStart - 2 * sectorsPerMB,
                                   sectorsPerMB) - 1)
      if osDataEnd // sectorsPerMB < SYSTEM_DISK_SIZE_SMALL_MB:
         # We expect a legacy layout system disk to have above 4GB of space
         # before the VMFS, with essential partitions and the scratch.
         raise RuntimeError("not enough disk space before VMFS datastore "
                            "(datastore starts at sector %u)" % vmfsStart)

      if disk.isUsb:
         disk.setPartition(7, LOCKER_FS, start, osDataEnd, LOCKER_LABEL)
      else:
         disk.setPartition(7, OSDATA_FS, start, osDataEnd, OSDATA_LABEL)

      # Preserve partition number of datastore if possible:
      # for GPT, make sure it is not a system-storage partition;
      # for MBR, the existing partition number should already be a
      #          primary number 2 or 3.
      dsPartNum = datastore.num
      if disk.isGpt or dsPartNum in (1, 5, 6, 7):
         dsPartNum = 8
      disk.setPartition(dsPartNum, FS_TYPE_VMFS, vmfsStart, vmfsEnd, vmfsLabel)
   else:
      if disk.isUsb:
         mediaEndMB = getMediaSize(sysMediaSize)
         if disk.sizeInMB > mediaEndMB:
            mediaEnd = calcEndSector(0, mediaEndMB, sectorsPerMB)
         else:
            mediaEnd = ...
         disk.setPartition(7, LOCKER_FS, start, mediaEnd, LOCKER_LABEL)
      else:
         partitionOsdataDisk(disk, partNum=7, startLba=start,
                             createDatastore=createDatastore,
                             sysMediaSize=sysMediaSize)

   # ESX always installs a GPT, except if upgrading from MBR, and a VMFS
   # datastore that must be preserved already occupies the last sectors of the
   # disk (prevents from writing the backup GPT).
   forceMbr = keepDatastore and not disk.isGpt

   # Apply partition table to disk and format them. Ignore errors when
   # formatting regular VMFS because the disk type might not support it,
   # e.g. LUNs that don't support ATS. This will ensure that system-storage
   # still comes up and the user can identify why a datastore couldn't
   # be formatted from a sysAlert in esxdisk.autoFormat().
   disk.syncPartitions(autoFormat=True, keepDatastore=keepDatastore,
                       forceMbr=forceMbr, ignoreVmfsFormatError=True)

def convertScratchDisk(disk, scratchPartNum):
   """Convert the disk containing a legacy scratch into an OSData volume.

   If there is a coredump partition immediately following the scratch partition,
   then also consolidate it with the scratch into a single OSData volume.
   Preserve existing datastore.

   @param disk: EsxDisk object with partition table scanned already
   @param scratchPartNum: scratch partition number
   """
   diskParts = disk.getPartitionsSortedByStart()
   scratchPart = None
   scratchIndex = 0
   corePart = None
   coreIndex = 0
   otherPart = None

   for index, (partNum, part) in enumerate(diskParts):
      if part.num == scratchPartNum:
         scratchIndex = index
         scratchPart = part
      elif part.fsType == FS_TYPE_VMKCORE:
         coreIndex = index
         corePart = part
      else:
         otherPart = part

   if scratchPart is None:
      raise FileNotFoundError("%s scratch partition #%u does not exist" %
                              (disk.name, scratchPartNum))

   # If only scratch (plus coredump) partition(s) exist on the disk, just
   # repartition the entire disk with only OSDATA.
   if disk.numPartitions <= 2 and otherPart is None:
      disk.clearPartitions()
      partitionOsdataDisk(disk, doPartSync=True)
      return

   if corePart is not None and abs(coreIndex - scratchIndex) == 1:
      # Consolidate if scratch and coredump partitions are contiguous.
      start = min(scratchPart.start, corePart.start)
      end = max(scratchPart.end, corePart.end)
      disk.removePartition(corePart.num)
   else:
      start = scratchPart.start
      end = scratchPart.end

   disk.setPartition(scratchPartNum, OSDATA_FS, start, end, OSDATA_LABEL)
   disk.syncPartitions(autoFormat=True, keepDatastore=True)


def partitionOsdataDisk(disk, partNum=7, startLba=None, createDatastore=False,
                        doPartSync=False, sysMediaSize=None):
   """Create the OSDATA and datastore (optional) partitions.

   The provided disk is assumed to be OSDATA-compatible (high-endurance).

   Params:
     partNum: partition number to assign to OSData.
     startLba: starting LBA of OSData; if None then OSData is set at first
               usable 1MB aligned sector.
     createDatastore: create a new VMFS volume on remaining space.
     doPartSync: sync partition table to disk.
     sysMediaSize: System-storage media size option (one of MEDIA_SIZES key)
                   to limit the (supported) size of OSData.
   """
   if not disk.supportsOsdata:
      raise RuntimeError('%s: disk device does not support OSDATA' % disk.name)

   sectorsPerMB = 1024 * 1024 // disk.sectorSize
   if startLba is None:
      startLba = sectorsPerMB

   # Create a datastore only on non-USB that is large enough.
   createDatastore = createDatastore and disk.supportsDatastore

   # Use vmkernel boot option for OSData size if one is supplied.
   # Prioritize systemMediaSize followed by autoPartitionOSDataSize.
   bootOpts = getVmkBootOptions()
   if SYSTEM_MEDIA_SIZE_OPT in bootOpts:
      osDataEndMB = getMediaSize(sysMediaSize)
   else:
      try:
         osDataSizeOpt = int(bootOpts[AUTOPART_OSDATA_SIZE_OPT])
         if osDataSizeOpt == 0:
            osDataEndMB = OSDATA_MAX_END_MB
         else:
            if osDataSizeOpt < VMFSL_MIN_SIZE_MB:
               sysAlert("OSData size %d is too small (minimum allowed %d)" %
                        (osDataSizeOpt, VMFSL_MIN_SIZE_MB))
            assert osDataSizeOpt >= VMFSL_MIN_SIZE_MB

            osDataEndMB = startLba // sectorsPerMB + osDataSizeOpt
      except AssertionError:
         raise
      except Exception:
         # Set OSData size based on system-storage media size option
         osDataEndMB = getMediaSize(sysMediaSize)

   if disk.sizeInMB > osDataEndMB:
      # For media > 128GB, only use up to 128GB for system storage
      osDataEnd = calcEndSector(0, osDataEndMB, sectorsPerMB)
      if createDatastore and (disk.sizeInMB - osDataEndMB) > VMFS_MIN_SIZE_MB:
         # Create a datastore for non-USB disk if the caller asks to create one.
         vmfsStart = roundupMBSector(osDataEnd + sectorsPerMB, sectorsPerMB)
         vmfsEnd = ...
         vmfsLabel = getNextDatastoreLabel()
      else:
         vmfsStart = vmfsEnd = None
   else:
      # OSdata occupies the whole disk.
      osDataEnd = ...
      vmfsStart = vmfsEnd = None

   disk.setPartition(partNum, OSDATA_FS, startLba, osDataEnd, OSDATA_LABEL)
   if vmfsStart is not None:
      disk.setPartition(8, FS_TYPE_VMFS, vmfsStart, vmfsEnd, vmfsLabel)

   if doPartSync:
      disk.syncPartitions(autoFormat=True, keepDatastore=not createDatastore,
                          ignoreVmfsFormatError=True)

def createDatastore(disk, label=None):
   """Create a new VMFS datastore on the entire disk.
   """
   assert disk.isEmpty and disk.supportsDatastore

   if label is None:
      label = getNextDatastoreLabel()

   # Set the first LBA at 2MB offset, which provides some buffer space at the
   # start of the disk.
   sectorsPerMB = 1024 * 1024 // disk.sectorSize
   startLba = 2 * sectorsPerMB

   disk.setPartition(1, FS_TYPE_VMFS, startLba, ..., label)
   disk.syncPartitions(autoFormat=True)
