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

"""Low-level operations related to ESX disks and partitions.
"""
import os
import time
import logging

from systemStorage import *

if IS_ESX:
   import json
   import re
   from esxutils import runCli, sysAlert
   from systemStorage.bootbank import scanEsxOnDisks
   from systemStorage.esxfs import fsRescan
   from systemStorage.vsanDisk import getVsanDisks, isVsanPresent, VsanDisk
   from vmSyslogUtils import execCommand

from systemStorage.blockdev import BlockDev, Partition
from systemStorage.gpt import Gpt
from systemStorage.mbr import MbrPartitionTable
from systemStorage.vfat import isVfatPartition, mkfsvfat

DD_GEOMETRY_SECTOR_SIZE = 512
DD_GEOMETRY_NUM_HEADS = 64
DD_GEOMETRY_SECTORS_PER_TRACK = 32

# Disk filters.
LOCAL_FILTER = 'local'
REMOTE_FILTER = 'remote'
USB_FILTER = 'usb'
ESX_FILTER = 'esx'
LOCALESX_FILTER = 'localesx'
REMOTEESX_FILTER = 'remoteesx'
NONSSD_FILTER = 'nonssd' # Not for direct use in firstdisk.
ESX_FILTERS = (ESX_FILTER, LOCALESX_FILTER, REMOTEESX_FILTER)

def mkfs(volumePath, fsType, volumeName, numBlocks, uuid=None,
         isRegularFile=False, offset=0, sectorSize=512):
   """Format a partition.

   For dd-image partitions, we only support VFAT. When creating OSDATA,
   auto-append a UUID to its name.
   """
   if fsType in (FS_TYPE_UEFI_SYSTEM, FS_TYPE_VFAT):
      mkfsvfat(volumePath, sectorSize, sectorSize, numBlocks, volumeName,
               isRegularFile, offset, uuid)
   else:
      if not IS_ESX:
         raise RuntimeError("Cannot format a VMFS partition on non-ESX env")

      from systemStorage.vmfsl import createVMFS
      appendUUID = (fsType == FS_TYPE_VMFS_L and
                    volumeName in (OSDATA_LABEL, LOCKER_LABEL))
      createVMFS(volumePath, fsType, volumeName, appendUUID=appendUUID)

def execLocalcli(cmd):
   """Execute localcli command and return result in a dictionary.

   runCli/EsxcliPy is used first, but in case or an error, especially when
   32-bit IMA plugins causing the call to fail, fallback to localcli.
   This means less performance impact when everything is 64-bit.
   """
   try:
      return runCli(cmd, evalOutput=True)
   except Exception:
      _, out, _ = execCommand(['/bin/localcli', '--formatter=json'] + cmd)
      return json.loads(out)

def getStoragePaths():
   """Return a dict of all ESX storage paths, indexed by device name.
   """
   paths = execLocalcli(['storage', 'core', 'path', 'list'])
   return {path['Device']: path for path in paths}

def getDiskPathInfo(diskName):
   """Get storage path info for a disk.
   """
   return execLocalcli(['storage', 'core', 'path', 'list', '-d', diskName])[0]

def getPathAdapterName(pathInfo):
   """Return the adapter name of a device from path info.
   """
   return pathInfo['Adapter']

def getStorageAdapters():
   """Return a dict of all ESX storage adapters, indexed by adapter name.
   """
   adapters = execLocalcli(['storage', 'core', 'adapter', 'list'])
   return {adapter['HBA Name']: adapter for adapter in adapters}

def getDiskAdapterInfo(diskName=None, pathInfo=None):
   """Get the adapter info dict of a disk.

   With storage path of a disk, we can find the adapter name, and then adapter
   info can be fetched.
   """
   assert (diskName or pathInfo), 'Either diskName or pathInfo must be supplied'
   if pathInfo is None:
      pathInfo = getDiskPathInfo(diskName)
   adapterName = getPathAdapterName(pathInfo)
   adapters = getStorageAdapters()
   return adapters[adapterName]


class PartitionTable(object):
   """Class to manage a disk partition table.
   """
   GPT_PART_TABLE = 'gpt'
   MBR_PART_TABLE = 'mbr'

   def __init__(self, sectorSize=512):
      self.partitions = {}
      self._type = None
      self._bootPart = None

   def setPartition(self, partNum, fsType, start, end, label, guid=None,
                    bootable=False, uuid=None):
      """Add/modify a partition on this disk, and return the partition info.
      """
      part = Partition(partNum, fsType, start, end, label, guid=guid,
                       bootable=bootable, uuid=uuid)
      self.partitions[partNum] = part

      if bootable:
         self._bootPart = partNum

      return part

   def removePartition(self, partNum):
      """Remove a partition on this disk.
      """
      del(self.partitions[partNum])
      if partNum == self._bootPart:
         self._bootPart = None

   def scan(self, blockDev):
      """Read the partition table from disk.
      """
      try:
         # Try MBR first.
         mbr = MbrPartitionTable()
         mbr.scan(blockDev)
         self._type = self.MBR_PART_TABLE
      except ValueError:
         # Protective MBR or invalid MBR seen, try GPT.
         gpt = Gpt(blockDev.numSectors, blockDev.sectorSize)
         gpt.scan(blockDev.readBlock)
         self._type = self.GPT_PART_TABLE

      if self._type == self.MBR_PART_TABLE:
         for partNum, part in mbr.partitions.items():
            self.setPartition(partNum, part.fsType, part.start, part.end,
                              part.label, bootable=part.bootable)
      else:
         for partNum, part in gpt.partitions.items():
            self.setPartition(partNum, part.fsType, part.start, part.end,
                              part.label, guid=part.guid)

   def syncGpt(self, blockDev, skipBackupGpt=False, mbrBootCode=None):
      """Write the partition table to disk as a GPT.
      """
      gpt = Gpt(blockDev.numSectors, blockDev.sectorSize)

      if mbrBootCode is not None:
         gpt.mbrBootCode = mbrBootCode

      for partNum, part in self.partitions.items():
         gpt.setPartition(partNum, part.fsType, part.start, part.end,
                          part.label, guid=part.guid)

      if self._bootPart is not None:
         gpt.setBootPartition(self._bootPart)

      gpt.sync(blockDev.writeBlock, skipBackupGpt=skipBackupGpt)

   def syncMbr(self, blockDev, bootCode):
      """Write the partition table to disk as a legacy MBR.
      """
      mbr = MbrPartitionTable()
      for partNum, part in self.partitions.items():
         mbr.setPartition(partNum, part.fsType, part.start, part.end,
                          bootable=(partNum==self._bootPart))
      mbr.sync(blockDev, bootCode)

   def iterPartitions(self, orderByStart=False):
      """Iterate over this disk's partitions.

      Partitions are sorted by partition number, or by starting sector if
      orderByStart=True.
      """
      if orderByStart:
         return sorted(self.partitions.items(), key=lambda p: p[1].start)
      return sorted(self.partitions.items(), key=lambda p: p[0])


class _EsxDisk(BlockDev):
   """Class to manage partitions compliant with ESX 7.0+ layout.
   """

   def clearPartitions(self):
      """Clear the partition map in the PartitionTable object.

      This does not flush the partition table on the disk.
      """
      self._pt.partitions = {}
      self._pt._bootPart = None

   def setPartition(self, partNum, fsType, start, end, label, **kwargs):
      """Add/modify a partition on this disk.
      """
      if fsType == FS_TYPE_VMFS and self.isUsb:
         # USB storage devices have a limited lifetime, and are particularly
         # sensitive to writes. This is incompatible with write-intensive VMFS
         # datastores.
         sysAlert("VMFS partition detected on %s. VMFS is not supported on "
                  "USB devices and risks data corruption or lost." % self.path)

      if end is Ellipsis:
         end = Gpt(self.numSectors, self.sectorSize).lastUsableLba

      return self._pt.setPartition(partNum, fsType, start, end, label, **kwargs)

   def removePartition(self, partNum):
      """Remove a partition on this disk.
      """
      self._pt.removePartition(partNum)

   def _supportedFormatFsTypes(self, keepDatastore=False):
      """Return a list of supported filesystem types that can be formatted
      """
      fsTypes = [FS_TYPE_UEFI_SYSTEM, FS_TYPE_VFAT, FS_TYPE_VMFS_L]
      if not keepDatastore:
         fsTypes.append(FS_TYPE_VMFS)
      return fsTypes

   def _formatPartition(self, part, partNum):
      if self._isRegularFile:
         volumePath = self.path
         offset = part.start * self.sectorSize
      else:
         volumePath = "%s:%s" % (self.path, partNum)
         offset = 0

      mkfs(volumePath, part.fsType, part.label, part.numSectors,
           part.uuid, self._isRegularFile, offset, self.sectorSize)

   def erasePartitionMetadata(self, part, blocks=4):
      """Erase metadata blocks of the partition.
      """
      LVM_DEV_HEADER_OFFSET = MiB // self.sectorSize
      self.eraseBlocks(part.start, blocks)
      if part.fsType in (FS_TYPE_VMFS, FS_TYPE_VMFS_L):
         # Erase the LVM header of VMFS
         self.eraseBlocks(part.start + LVM_DEV_HEADER_OFFSET, blocks)

   def autoFormat(self, keepDatastore=False, ignoreVmfsFormatError=False):
      """Format all partitions on this disk.

      @keepDatastore: when set, do not auto format VMFS datastore.
      @ignoreVmfsFormatError: ignore errors when formatting a VMFS partition.
      """
      fsTypes = self._supportedFormatFsTypes(keepDatastore)

      for partNum, part in self._pt.iterPartitions():
         if part.fsType in fsTypes:
            try:
               self._formatPartition(part, partNum)
            except OSError as e:
               sysAlert("Failed to format %s partition %d as %s: %s" %
                         (self.path, partNum, part.fsType, e))
               if part.fsType != FS_TYPE_VMFS or not ignoreVmfsFormatError:
                  raise

   def formatPartition(self, partNum):
      """Format the specified partition on this disk.
      """
      for pnum, part in self._pt.iterPartitions():
         if pnum == partNum:
            self._formatPartition(part, partNum)
            break
      else:
         raise FileNotFoundError("Partition %u not found on disk %s" %
                                 (partNum, self.name))

   def scanPartitions(self):
      """Read the partition table from disk.
      """
      self.open(os.O_RDONLY)
      try:
         try:
            self._pt.scan(self)
         except ValueError:
            # no valid partition table found
            pass
      finally:
         self.close()

   def syncPartitions(self, autoFormat=False, keepDatastore=False,
                      skipBackupGpt=False, mbrBootCode=None, forceMbr=False,
                      ignoreVmfsFormatError=False):
      """Write the partition table to disk.

      This function overwrites any existing partitions. New partitions are
      formatted if @autoFormat=True.
      VMFS datastore formatting is skipped if @keepDatastore=True.
      Backup GPT at the end of the disk is skipped if @skipBackupGpt=True.
      @ignoreVmfsFormatError=True will ignore formatting errors on VMFS
      partitions.
      """
      fsTypes = self._supportedFormatFsTypes(keepDatastore)

      self.open(os.O_WRONLY)
      try:
         # Retry writing the partition table a few times because the
         # filesystems might not be fully unmounted with references to the
         # volumes still existing, leading to Read-Only errors.
         # The short wait time of 0.5s between intervals is because the
         # volume cache flushing out of these references is done in software
         # which should be relatively fast.
         retries = 10
         while retries > 0:
            try:
               if forceMbr:
                  self._pt.syncMbr(self, mbrBootCode)
               else:
                  self._pt.syncGpt(self, skipBackupGpt=skipBackupGpt,
                                   mbrBootCode=mbrBootCode)
               break
            except Exception as e:
               if retries == 1:
                  if IS_ESX:
                    sysAlert("Failed to update partition table due to: %s" % e)
                  raise
               else:
                  time.sleep(0.5)
            retries -= 1

         # Erase metadata so that fsRescan doesn't pick up stale partition info
         if autoFormat and isinstance(self, EsxDisk):
            for _, part in self._pt.iterPartitions():
               if part.fsType in fsTypes:
                  self.erasePartitionMetadata(part)
      finally:
         self.close()

      if isinstance(self, EsxDisk):
         # Rescan disks so new partitions appear under /dev/disks/.
         fsRescan()

      if autoFormat:
         self.autoFormat(keepDatastore=keepDatastore,
                         ignoreVmfsFormatError=ignoreVmfsFormatError)

   @property
   def isEmpty(self):
      """True if this disk contains no partitions.
      """
      return len(self._pt.partitions) == 0

   @property
   def isGpt(self):
      """True if this partition table is a GPT.
      """
      return self._pt._type == PartitionTable.GPT_PART_TABLE

   @property
   def hasBootbanks(self):
      """True if this disk contains two ESX bootbanks.
      """
      uefiSystem = 0
      vfat = 0

      for _, part in self._pt.iterPartitions():
         if part.fsType == BOOTPART_FS:
            uefiSystem += 1
         elif part.fsType == BOOTBANK_FS:
            vfat += 1

         if uefiSystem == 1 and vfat >= 2:
            # boot + bootbank1 + bootbank2 [+ legacyLocker + legacyScratch]
            return True

      return False

   @property
   def hasLegacyBootPart(self):
      """True if this disk contains a legacy boot partition (<= 4MiB).
      """
      LEGACY_BOOTPART_SIZE = 4 * MiB
      LEGACY_GPT_BOOTPART = 1
      LEGACY_MBR_BOOTPART = 4

      bootPartNum = LEGACY_GPT_BOOTPART if self.isGpt else LEGACY_MBR_BOOTPART

      for partNum, part in self._pt.iterPartitions():
         # Legacy MBR boot partition is of type FAT16.
         if (part.fsType in (BOOTPART_FS, FS_TYPE_FAT16) and
             partNum == bootPartNum and
             part.numSectors * self.sectorSize <= LEGACY_BOOTPART_SIZE):
            return True

      return False

   @property
   def hasLegacyScratch(self):
      """True if this disk contains a legacy VFAT scratch partition.
      """
      # Minimum size of a legacy Scratch partition throughout ESX versions.
      MIN_SCRATCH_PARTITION_SIZE = 4000000000

      for _, part in self._pt.iterPartitions():
         if part.fsType == LEGACY_SCRATCH_FS:
            partSize = part.numSectors * self.sectorSize
            if partSize >= MIN_SCRATCH_PARTITION_SIZE:
               # Partition type is VFAT and it is at least
               # MIN_SCRATCH_PARTITION_SIZE big.
               return True
      return False

   @property
   def hasOsdata(self):
      """True if this disk contains an OSDATA partition.
      """
      return self.supportsOsdata and self.hasSystemVmfsl

   @property
   def hasSystemVmfsl(self):
      """True if this disk contains a system VMFS-L partition.
      """
      from systemStorage.vmfsl import getSystemVmfslDesc

      for partNum, part in self._pt.iterPartitions():
         if part.fsType in (FS_TYPE_VMFS_L, FS_TYPE_VMFS):
            # A partition can convert between VMFS and VMFS-L.
            devPath = "%s:%d" % (self.path, partNum)
            try:
               getSystemVmfslDesc(devPath)
            except ValueError:
               # Not a system VMFS-L partition.
               continue
            else:
               return True

      return False

   @property
   def supportsDatastore(self):
      """True if this disk can have a VMFS datastore.
      """
      return not self.isUsb and self.sizeInMB >= VMFS_MIN_SIZE_MB

   @property
   def supportsOsdata(self):
      """True if this disk can have an OSDATA partition.
      """
      return not self.isUsb and self.sizeInMB >= VMFSL_MIN_SIZE_MB

   @property
   def numPartitions(self):
      """Get number of partitions
      """
      return len(self._pt.partitions)

   def getPartition(self, partNum):
      """Get partition by number.
      """
      return self._pt.partitions[partNum]

   def getPartitionsByFsType(self, fsType):
      """Get a list of partition by fsType.
      """
      return self.getPartitionsByFsTypes([fsType])

   def getPartitionsByFsTypes(self, fsTypes):
      """Get a list of partition by a list of fsTypes.
      """
      return [part for _, part in self._pt.iterPartitions()
              if part.fsType in fsTypes]

   def getPartitionsSortedByStart(self):
      """Return a list of disk partitions sorted by starting sector.
      """
      return self._pt.iterPartitions(orderByStart=True)

   def getBootPartition(self):
      """Get the ESX boot partition on this disk.
      """
      parts = self.getPartitionsByFsType(BOOTPART_FS)
      numBootParts = len(parts)
      if numBootParts == 0:
         raise RuntimeError("%s: no boot partition found" % self.path)
      elif numBootParts > 1:
         raise RuntimeError("%s: more than one boot partition found (%u)" %
                            (self.path, numBootParts))
      return parts[0]

   def getDatastore(self):
      """Retrieve info about the datastore partition on this disk (if any).
      """
      partitions = self.getPartitionsByFsType(FS_TYPE_VMFS)
      numDatastores = len(partitions)
      assert numDatastores <= 1, ("%s: %u datastores found on disk (maximum "
                                  "allowed is 1)" % (self.name, numDatastores))
      try:
         return partitions[0]
      except IndexError:
         return None

   @property
   def hasDatastore(self):
      """True if this disk contains a VMFS datastore.
      """
      return self.getDatastore() is not None


class EsxDD(_EsxDisk):
   """Class to manage ESX GPT partitioning on a regular file.
   """

   def __init__(self, filePath):
      DD_FILE_SIZE_IN_BYTES = 1 * GiB + 128 * MiB
      # XXX: Knowing the size of the disk is needed for the protective MBR
      # metadata among other things. We don't know in advance the size of the
      # disk so we use a size of 1152MB to at least accomodate the size of the
      # system boot partition + bootbank1 + bootbank2 + some padding.
      numSectors = DD_FILE_SIZE_IN_BYTES // DD_GEOMETRY_SECTOR_SIZE

      super().__init__(filePath,
                       DD_GEOMETRY_SECTOR_SIZE,
                       numSectors=numSectors,
                       numHeads=DD_GEOMETRY_NUM_HEADS,
                       sectorsPerTrack=DD_GEOMETRY_SECTORS_PER_TRACK,
                       isRegularFile=True)

      self._pt = PartitionTable()


class EsxDisk(_EsxDisk):
   """Class to access disk info and blocks on ESX.
   """

   def __init__(self, name, devInfo=None, vsanInfo=None, pathInfo=None,
                adapterInfo=None, devInfoOnly=False):
      if devInfo is None:
         self._info = execLocalcli(['storage', 'core', 'device', 'list',
                                    '--exclude-offline',
                                    '--device', name])[0]
      else:
         self._info = devInfo

      self.name = self._info['Device']
      self._adapterInfo = adapterInfo
      self._pathIds = None

      if devInfoOnly:
         # Only interested in device info
         return

      if pathInfo is None:
         pathInfo = getDiskPathInfo(name)
      self._pathInfo = pathInfo

      # Empty dict means the disk is already confirmed not to be in vSAN.
      if vsanInfo != {} and isVsanPresent():
         try:
            self._vsanDisk = VsanDisk(self.name, vsanInfo)
         except ValueError:
            # Not a vSAN disk.
            self._vsanDisk = None
      else:
         self._vsanDisk = None

      capacity = execLocalcli(['storage', 'core', 'device', 'capacity',
                               'list', '--device', self.name])[0]
      assert capacity['Device'] == self.name
      sectorSize = capacity['Logical Blocksize']

      if capacity['Format Type'].lower().startswith('4kn'):
         # Override sector size if 4Kn device, as storage layer also
         # provides software emulation of 512, but logical block size
         # should be 4k, not 512.
         sectorSize = 4096

      # Use stat to determine real disk size because 'esxcli storage core'
      # might return a smaller size than the actual disk size.
      st = os.stat(os.path.join(os.path.sep, 'dev', 'disks', name))
      sizeInMB = st.st_size // MiB
      numSectors = st.st_size // sectorSize

      super().__init__(self._info['Devfs Path'], sectorSize, numSectors)
      assert self.sizeInMB == sizeInMB

      self._pt = PartitionTable()

   def __str__(self):
       return '%s (disk %s) -- %s %s (%u MiB, %s)' % (
                  self.name,
                  self.path,
                  self.vendor,
                  self.model,
                  self.sizeInMB,
                  self.driverName)

   @property
   def devPath(self):
      """Full device path.
      """
      return os.path.join(os.path.sep, 'dev', 'disks', self.name)

   def isBackupGptHealthy(self):
      """Check if the backup GPT is present and healthy.
      """
      assert self.isGpt
      gpt = Gpt(self.numSectors, self.sectorSize)

      self.open(os.O_RDONLY)

      try:
         gpt.scan(self.readBlock)
         isHealthy = gpt.isBackupGptHealthy(self.readBlock)
      finally:
         self.close()

      return isHealthy

   @classmethod
   def fromDevInfo(cls, devInfo, vsanInfo, pathInfo, adapterInfo):
      """Instanciate an EsxDisk from the following device info dicts:

      devInfo      - dict that contains disk info.
      vsanInfo     - dict that contains vSAN disk info, can be an empty dict
                     indicating that the disk is not in vSAN.
      pathInfo     - dict that contains info for the disk's storage path.
      adapterInfo  - dict that contains info for the disk's storage adapter.
      """
      return cls(devInfo['Device'], devInfo, vsanInfo, pathInfo, adapterInfo)

   @property
   def isSsd(self):
      """True for SSD disks.
      """
      return self._info['Is SSD']

   @property
   def isUsb(self):
      """True for USB adapters.
      """
      return self._info['Is USB']

   @property
   def isLocal(self):
      """True for local disks.
      """
      return self._info['Is Local'] or self._info['Is Local SAS Device']

   @property
   def isFcSan(self):
      """True for SAN disks that are connected with FC.
      """
      return self._pathInfo['Transport'] == 'fc'

   @property
   def isActiveBootDisk(self):
      """True if ESX was booted from this disk.
      """
      return self._info['Is Boot Device']

   @property
   def isRemovable(self):
      """True for removable device.
      """
      return self._info['Is Removable']

   @property
   def vsanGroupUuid(self):
      """VSAN group UUID, or None if the disk does not belong to VSAN.
      """
      return None if self._vsanDisk is None else self._vsanDisk.diskGroupUuid

   @property
   def isVsanClaimed(self):
      """True if the disk belongs to a vSAN disk group.
      """
      return self._vsanDisk is not None

   @property
   def vendor(self):
      """The vendor of the disk.
      """
      return self._info['Vendor']

   @property
   def model(self):
      """The model of the disk.
      """
      return self._info['Model']

   @property
   def otherUids(self):
      """Other UIDs of the disk.
      """
      return self._info['Other UIDs']

   @property
   def deviceType(self):
      """The physical device of disk.
      """
      if self._info['Device Type'].strip() == 'CD-ROM':
         devType = 'cdrom'
      elif self._info['Is USB']:
         devType = 'usb'
      elif self._info['Is SSD']:
         devType = 'ssd'
      elif self._info['Is SAS']:
         devType = 'sas'
      elif self._info['Is Local']:
         devType = 'hdd'
      else:
         devType = 'unknown'
      return devType

   @property
   def driverName(self):
      """Driver name of this disk's adapter.
      """
      if self._pathInfo is None:
         self._pathInfo = getDiskPathInfo(self.name)
      if self._adapterInfo is None:
         self._adapterInfo = getDiskAdapterInfo(pathInfo=self._pathInfo)
      return self._adapterInfo['Driver']

   @property
   def interfaceType(self):
      """Interface type (transport type) of the disk's storage path.
      """
      if self._pathInfo is None:
         self._pathInfo = getDiskPathInfo(self.name)
      return self._pathInfo['Transport']

   @property
   def pathIds(self):
      """Disk storage adapter info tuple.

      Return the storage adapter info as tupple which can be used (by the ESX
      installer) to identify disks in a determistic way.

      Tuple format: ((name, number), channel #, adapter #, LUN #)
      e.g.          mpx.vmhba1:C0:T2:L0 will have (('vmhba', 1), 0, 2, 0)
      """
      if self._pathIds is None:
         if self._pathInfo is None:
            self._pathInfo = getDiskPathInfo(self.name)
         adapterName = getPathAdapterName(self._pathInfo)
         parts = [p for p in re.split('^(\D+)(\d+)$', adapterName) if p]
         adapterTuple = ((parts[0], int(parts[1])) if len(parts) == 2
                         else (adapterName,))
         self._pathIds = (adapterTuple, self._pathInfo['Channel'],
                          self._pathInfo['Target'], self._pathInfo['LUN'])
      return self._pathIds

   def getVfatPartitions(self):
      """Scan through the disk for all VMware created VFAT partitions.

      This is done by inspecting the second sector of the partition for the
      VMware FAT16 magic string.
      """
      vfatTypes = (FS_TYPE_VFAT, FS_TYPE_UEFI_SYSTEM, FS_TYPE_FAT16)
      partitions = []
      for partNum, part in self._pt.iterPartitions():
         if (part.fsType in vfatTypes and
             isVfatPartition("%s:%d" % (self.path, partNum), self.sectorSize)):
            partitions.append(partNum)
      return partitions


def filterDisks(disks, filterName, esxScanDict=None):
   """Filter a list of disks with one filter.

   ESXi scan result is required if one of the esx filters is included.
   A ESXi scan dict can be supplied to avoid a rescan.
   Supported fitler names:
      * local     - match a local disk.
      * remote    - match a remote disk.
      * usb       - match a usb disk.
      * nonssd    - match a disk that is not ssd (most likely HDD or USB);
                    not an installer 'firstdisk' filter, but useful the for
                    the --ignoressd kickstart option.
      * esx       - match a disk with ESXi installed;
                    a local ESXi disk will be returned before a remote one.
      * localesx  - match a local ESXi disk.
      * remoteesx - match a remote ESXi disk.
      * <other>   - other values will be treated as vendor, model or
                    driverName of the disk.
   """
   local = lambda d: d.isLocal
   remote = lambda d: not d.isLocal and not d.isUsb
   usb = lambda d: d.isUsb
   nonssd = lambda d: not d.isSsd
   localEsx = lambda d, esxInfo: d.name in esxInfo and d.isLocal
   remoteEsx = lambda d, esxInfo: (d.name in esxInfo and not d.isLocal
                                   and not d.isUsb)
   diskAttr = lambda d, attr: attr.strip().lower() in (
                                                 d.vendor.strip().lower(),
                                                 d.model.strip().lower(),
                                                 d.driverName.strip().lower())

   simpleFilters = {LOCAL_FILTER: local,
                    REMOTE_FILTER: remote,
                    USB_FILTER: usb,
                    NONSSD_FILTER: nonssd}

   if filterName in simpleFilters:
      filterFunc = simpleFilters[filterName]
      result = [d for d in disks if filterFunc(d)]
   elif filterName in ESX_FILTERS:
      # Lazy initiate ESXi info when required.
      if esxScanDict is None:
         diskNames = [d.name for d in disks]
         esxScanDict = scanEsxOnDisks(diskNames)
      if filterName == ESX_FILTER:
         # Local esx first and then remote esx,
         result = [d for d in disks if localEsx(d, esxScanDict)]
         result += [d for d in disks if remoteEsx(d, esxScanDict)]
      else:
         filterFunc = localEsx if filterName == LOCALESX_FILTER else remoteEsx
         result = [d for d in disks if filterFunc(d, esxScanDict)]
   else:
      # Treat as one of disk's vendor, model and driverName.
      result = [d for d in disks if diskAttr(d, filterName)]
   return result

def getFirstDiskOrder(disks, filters, esxScanDict=None, sortByPathIds=True,
                      skipSsd=False, skipNonEsx=False, checkEligible=False):
   """Return a list of firstdisk choices according to a list of filters.

   A disk will appear early if it satisfies a filter that is listed early,
   and will not be listed twice.
   Parameters:
      disks         - a list of EsxDisk objects.
      filters       - a list of firstdisk filter names.
      esxScanDict   - result dict of ESXi scanning.
      sortByPathIds - sort all disks by their path IDs, this makes
                      sure a predictable result is returned.
      skipSsd       - set to True to return only a HDD or a USB.
      skipNonEsx    - set to True to return only a disk that has ESXi
                      (firstdisk for upgrade).
      checkEligible - set to True to return only a disk that is large
                      enough (eligible) for ESXi installation.
   """
   esxScanDict = None
   if checkEligible:
      disks = [d for d in disks if d.sizeInMB >= SYSTEM_DISK_SIZE_SMALL_MB]
   if skipSsd:
      disks = filterDisks(disks, NONSSD_FILTER)
   if skipNonEsx:
      esxScanDict = scanEsxOnDisks([d.name for d in disks])
      disks = filterDisks(disks, ESX_FILTER, esxScanDict=esxScanDict)
   if sortByPathIds:
      disks = sorted(disks, key=lambda d: d.pathIds)

   diskNames = set()
   diskList = []
   if disks:
      for filterName in filters:
         if filterName in ESX_FILTERS and esxScanDict is None:
            # Reuse scanning result for multiple ESX fiters.
            esxScanDict = scanEsxOnDisks([d.name for d in disks])
         result = filterDisks(disks, filterName, esxScanDict=esxScanDict)
         for disk in result:
            if disk.name not in diskNames:
               diskList.append(disk)
               diskNames.add(disk.name)
   return diskList

def getFirstDisk(disks, filters, esxScanDict=None, sortByPathIds=True,
                 skipSsd=False, skipNonEsx=False, checkEligible=False):
   """Similar to getFirstDiskOrder(), but only return one disk.

   None is returned if there is no suitable firstdisk.
   """
   diskList = getFirstDiskOrder(disks, filters, esxScanDict, sortByPathIds,
                                skipSsd, skipNonEsx, checkEligible)
   return diskList[0] if diskList else None

def getDiskCollection():
   """Get a dict of all disks in the system, indexed by name.
   """
   return {disk.name: disk for disk in iterDisks()}

def iterDisks():
   """Iterate over ESX disks.
   """
   devices = execLocalcli(['storage', 'core', 'device', 'list',
                           '--exclude-offline'])

   # Get a dictionary of all vSAN disks, indexed by disk name.
   # When vSAN is not present on the system, just use an empty dict.
   vsanDisks = getVsanDisks() if isVsanPresent() else dict()

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

   # Get dictionaries for storage paths and adapters.
   try:
      pathDict = getStoragePaths()
      adapterDict = getStorageAdapters()
   except Exception as e:
      logger.warning("Error querying storage paths and adapters: %s", e)
      return

   for dev in devices:
      if dev['Device Type'].strip() == 'CD-ROM':
         continue

      if dev['Size'] == 0:
         # Skip zero-sized device, which has access problems and can't be used.
         continue

      if dev['Is Pseudo']:
         # Skip pseudo devices (which are not actual disks, but management
         # interfaces to send commands to the storage array controller).
         continue

      if dev['Is Offline']:
         # Devices which are offline, or turned off by the administrator
         # should not be enumerated.
         continue

      if dev['Status'].strip() == 'off':
         continue

      devName = dev['Device']

      # Assign an empty dict when the disk is not in vSAN.
      vsanInfo = vsanDisks.get(devName, dict())

      # Get path/adapter info.
      try:
         pathInfo = pathDict[devName]
      except KeyError:
         # pathDict may incl. only a subset of the devices. A path can be
         # still in the initializing state or the disk is currently removed
         # (PR2608864).
         continue

      adapterInfo = adapterDict[getPathAdapterName(pathInfo)]

      try:
         yield EsxDisk.fromDevInfo(dev, vsanInfo, pathInfo, adapterInfo)
      except Exception as e:
         logger.warning("Failed to initialize storage device (%s) info: %s",
                        devName, e)
