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

"""ESX coredump utility functions.
"""
import os

from advcfg import getVmkCfgOpt
from esxutils import runCli, getVmkBootOptions, EsxcliError
import vmkctl

MB = 1024 * 1024

class CoreDumpLocation:
   """Class for coredump metadata (activated, type, name, size)
   """
   def __init__(self, active, configured, coreType, corePath, coreName='',
                coreSize=0):
      self.active = active
      self.configured = configured
      self.type = coreType
      self.path = corePath
      self.name = coreName
      self.size = coreSize

   @property
   def isPartition(self):
      return self.type == 'partition'

   @property
   def isFile(self):
      return self.type == 'file'

   @property
   def isNetwork(self):
      return self.type == 'network'

   @property
   def exists(self):
      if self.isNetwork:
         return True
      return os.path.exists(self.path)

   def activate(self):
      """Set this coredump as the active one
      """
      if self.isFile:
         enableCoredumpToFile(self.path)
      elif self.isPartition:
         enableCoredumpToPart(self.name)
      else:
         enableCoredumpToNetwork()

def autoCoredumpEnabled():
   """Check if autoCreateDumpFile is enabled
   """
   return getVmkCfgOpt('autoCreateDumpFile')

def getDiskDumpSlotSize():
   """Get the configured diskDumpSlotSize, the amount of storage available in
      core-partition for the system to use.
   """
   try:
      return int(getVmkBootOptions()['diskDumpSlotSize'])
   except Exception:
      return getVmkCfgOpt('diskDumpSlotSize')

def getCoredumpPartition(configuredOnly=True):
   """Return the current coredump partition.

      If configuredOnly=True, then return the configured device name if it
         exists, otherwise none;
      Else return the whole partition object as provided from localcli.
   """
   out = runCli(['system', 'coredump', 'partition', 'get'], evalOutput=True)
   if not configuredOnly:
      return out

   if not out['Configured']:
      return None
   return out['Configured']

def unconfigureCoredumpToPart():
   """Unconfigure coredump to partition.
   """
   runCli(['system', 'coredump', 'partition', 'set', '-u'])

def enableCoredumpToPart(partDev=None):
   """Enable coredump to partition with auto config.
   """
   if partDev is None:
      runCli(['system', 'coredump', 'partition', 'set', '-e', 'true', '-s'])
   else:
      runCli(['system', 'coredump', 'partition', 'set', '-p', partDev])
      runCli(['system', 'coredump', 'partition', 'set', '-e', 'true'])

def unconfigureCoredumpToFile():
   """Unconfigure coredump to file.
   """
   runCli(['system', 'coredump', 'file', 'set', '-u'])

def enableCoredumpToFile(filePath=None):
   """Enable coredump to file with auto config.
   """
   if filePath is None:
      runCli(['system', 'coredump', 'file', 'add', '-a', '-e', 'true'])
      runCli(['system', 'coredump', 'file', 'set', '-e', 'true', '-s'])
   else:
      runCli(['system', 'coredump', 'file', 'set', '-p', filePath])
      runCli(['system', 'coredump', 'file', 'set', '-e', 'true'])

def enableCoredumpToNetwork():
   """Enable network coredump as configured.
   """
   runCli(['system', 'coredump', 'network', 'set', '-e', 'true'])

def getCoredumpFile():
   """Get the active/configured coredump file
   """
   return runCli(['system', 'coredump', 'file', 'get'], evalOutput=True)

def getCoredumpFiles():
   """Identify the file(s) used for coredumps
   """
   return runCli(['system', 'coredump', 'file', 'list'], evalOutput=True)

def iterateCoredumpLocations(configuredOnly=False):
   """Iterate through the configured coredump locations on the system.
      If configuredOnly=True, then only retrieve configured core locations.
   """
   try:
      netConf = runCli(['system', 'coredump', 'network', 'get'],
                       evalOutput=True)
      yield CoreDumpLocation(netConf['Enabled'], netConf['Enabled'], 'network',
                             "%s:%u" % (netConf['Network Server IP'],
                             netConf['Network Server Port']))
   except Exception:
      pass

   try:
      if configuredOnly:
         corefile = getCoredumpFile()
         if corefile['Configured'] != '':
            if os.path.exists(corefile['Configured']):
               coreSize = os.path.getsize(corefile['Configured']) // MB
            else:
               coreSize = 0
            yield CoreDumpLocation(corefile['Active'] != '', True, 'file',
                                   corefile['Configured'], coreSize=coreSize)
      else:
         filesConf = getCoredumpFiles()

         # Coredump file list will only output files that actually exist
         # so there is no need to verify their existence.
         for f in filesConf:
            yield CoreDumpLocation(f['Active'], f['Configured'], 'file',
                                   f['Path'], coreSize=f['Size'] // MB)
   except Exception:
      pass

   try:
      if configuredOnly:
         part = getCoredumpPartition(False)
         name = part['Configured']
         if name != '':
            path = os.path.join(os.path.sep, "vmfs", "devices", "disks", name)
            if os.path.exists(path):
               partSizeMB = os.path.getsize(path) // MB
            else:
               partSizeMB = 0
            yield CoreDumpLocation(part['Active'] != '', True, 'partition',
                                   path, coreName=name, coreSize=partSizeMB)
      else:
         partsConf = runCli(['system', 'coredump', 'partition', 'list'],
                            evalOutput=True)
         # Partition list will only output partitions that actually exist
         # so there is no need to verify their existence.
         for p in partsConf:
            partSizeMB = os.path.getsize(p['Path']) // MB
            yield CoreDumpLocation(p['Active'], p['Configured'], 'partition',
                                   p['Path'], p['Name'], partSizeMB)
   except Exception:
      pass


def calcNeededCoredumpSize():
   """Calculate the coredump size recommended for this system
   """
   storage = vmkctl.StorageInfoImpl()
   coreSize = storage.GetRecommendedDumpSize() // MB
   return coreSize

def calcCoreFileSize(volPath, defaultSize=0, existingCoreSize=0, fitInVol=True):
   """Calculate the size of the coredump file for the VMFS volume.
   In addition to the defaultSize parameter, the dumpSize boot config
   option is also considered as part of the calculation, such that
   the largest core-size that can fit within the free-space is used as the
   final size.

   @param defaultSize      Size in MB recommended for coredump file. If 0 is
                           specified then the size is automatically calculated
                           as appropriate for the system.
   @param existingCoreSize Size in MB of existing coredump file in volume.
   @param fitInVol         Calculate size of core to fit within volume if
                           autosize and override sizes are too big.
   """
   if defaultSize == 0:
      defaultSize = calcNeededCoredumpSize()

   try:
      dumpSizeOption = int(getVmkBootOptions()['dumpSize'])
      if dumpSizeOption == 0:
         dumpSizeOption = defaultSize
   except Exception:
      dumpSizeOption = defaultSize

   vmfsUuid = os.path.basename(volPath)
   stat = os.statvfs(volPath)
   freeSpace = existingCoreSize + (stat.f_bavail * stat.f_frsize) // MB
   if freeSpace < min(dumpSizeOption, defaultSize) * 2:
      if fitInVol:
         # Both calculated and overriding sizes are too big to fit;
         # resize to half available space so that it can be created.
         # Subtract 1MB so that the underlying localcli coredump command
         # doesn't complain about insufficient space.
         defaultSize = freeSpace // 2 - 1
      else:
         raise RuntimeError('%s: insufficient space for coredump; '
                            'needed %uMB, free %uMB' %
                            (vmfsUuid, min(dumpSizeOption, defaultSize) * 2,
                             freeSpace))
   elif dumpSizeOption > defaultSize:
      # Ensure overriding size fits in freeSpace
      if freeSpace > dumpSizeOption * 2:
         defaultSize = dumpSizeOption
      elif fitInVol:
         # Can only fit up to half of free space (less 1MB)
         defaultSize = freeSpace // 2 - 1
      else:
         raise RuntimeError('%s: insufficient space for coredump; '
                            'needed %uMB, free %uMB' %
                            (vmfsUuid, dumpSizeOption * 2, freeSpace))

   return defaultSize


def assignCoreToFile(vmfsUuid, sizeInMB=0):
   """Assign the coredump-to-file location to the specified VMFS volume.

   Returns the path to the coredump file on success
   """
   volUuid = str(vmfsUuid)

   coreVol = os.path.join("/vmfs/volumes", volUuid)
   dumpDir = os.path.join(coreVol, "vmkdump")
   try:
      coreFiles = os.listdir(dumpDir)
   except Exception:
      coreFiles = []

   coreFilePath = None
   coreSizeMB = 0

   # Find largest core file
   for coreFile in coreFiles:
      if coreFile.endswith('.dumpfile'):
         filepath = os.path.join(dumpDir, coreFile)
         st = os.stat(filepath)
         size = st.st_size // MB
         if size > coreSizeMB:
            coreFilePath = filepath
            coreSizeMB = size

   if coreFilePath is None:
      # No coredump exists in the volume, so derive the size for it.
      sizeInMB = calcCoreFileSize(coreVol, sizeInMB)
   else:
      try:
         sizeInMB = calcCoreFileSize(coreVol, sizeInMB, coreSizeMB,
                                     fitInVol=False)
      except RuntimeError:
         # Insufficient space when attempting to upsize a coredump file.
         # Don't remove the coredump file and activate it below.
         sizeInMB = 0

      if coreSizeMB >= sizeInMB:
         enableCoredumpToFile(coreFilePath)
         return coreFilePath

      # Delete it so that it will be recreated
      runCli(['system', 'coredump', 'file', 'set', '--enable', 'False'])
      runCli(['system', 'coredump', 'file', 'remove', '--file', coreFilePath])

   # Create dumpfile at /vmfs/volumes/UUID/vmkdump/FILENAME.dumpfile

   os.makedirs(dumpDir, exist_ok=True)
   try:
      runCli(['system', 'coredump', 'file', 'add', '--datastore', volUuid,
              '--size', str(sizeInMB), '--enable', '1'])
   except EsxcliError as e:
      # Ignore already exists failure. This can happen when RAM or dumpSize
      # increased and no space available, so the corefile location was
      # skipped above.
      if 'already exists' not in str(e):
         raise

   # Verify dump file configuration
   try:
      configuredDumpFile = getCoredumpFile()

      # The coredump creation process above should've activated the file
      assert configuredDumpFile['Active'] != '', \
             "Coredump file in volume %s not active" % vmfsUuid

      return configuredDumpFile['Active']
   except Exception as e:
      raise OSError("%s: Coredump file assignment to volume failed: %s" %
                    (vmfsUuid, e))

