# Copyright 2014-2021 VMware, Inc.
# All rights reserved. -- VMware Confidential

"""Common ESX utils
"""
import logging
import os
import re
import shutil
import subprocess
import sys
import time
import uuid

from vmware import vsi

VMKFSTOOLS = '/sbin/vmkfstools'
SCRATCH_PARTITION = '/scratch'

def sysAbort(msg):
   """Panics the system with a message.
   """
   SYS_ABORT_NODE = '/system/sysAbort'
   vsi.set(SYS_ABORT_NODE, msg)

def sysAlert(msg):
   """Print a SysAlert message in vmkernel.log.
   """
   vsi.set('/system/sysAlert', msg)

class SysAlertHandler(logging.Handler):
   """Logging handler class to generate sys alerts
   """
   def emit(self, record):
      """Generate sysalert from log."""
      try:
         msg = self.format(record)
         sysAlert(msg)
      except (KeyboardInterrupt, SystemExit):
         raise
      except:
         self.handleError(record)

class EsxBootOpts(dict):
   """Class to retrieve ESX boot options
   """
   def __init__(self):
      opts = getVmkBootOptions()
      self.update(opts)

   def getBool(self, name, default=False):
      """Return a boolean ESX boot option.
      """
      assert default is None or default is True or default is False

      if default is None:
         default = False
      val = self.get(name, default)
      if val is None:
         # Only name key exists, treat as True
         return True

      if isinstance(val, bool):
         return val

      val = val.lower()
      if val not in ('1', 'true', '0', 'false'):
         raise ValueError("%s: vmkernel boot option is not a boolean (%s)" %
                          (name, val))
      return val in ('true', '1')

def isESX():
   """Return True if we're running on ESX.
   """
   return os.uname()[0] == 'VMkernel'

if isESX():
   import esxclipy

def isEsxInAVm():
   """Check whether ESX is running in a VM.
   """
   cpuInfo = vsi.get('/hardware/cpu/cpuInfo')
   if cpuInfo.get('nvoa', 0) == 1:
      dmiInfo = vsi.get('/hardware/bios/dmiInfo')
      if dmiInfo.get('vendorName', '') == 'VMware, Inc.':
         return True
   return False

class EsxcliError(Exception):
   pass

class NfsAlreadyExportedError(Exception):
   def __init__(self, volumeName):
      self.volumeName = volumeName
   def __str__(self):
      return "This NFS is already mounted as %s" % self.volumeName


class VmkConfigOpt(object):
   """Class to read/write VMKernel Config Options.
   """

   def __init__(self, group, name, strOpt=False):
      super().__init__()
      self._strOpt = strOpt
      optType = 'strOpts' if strOpt else 'intOpts'
      self._path = '/config/%s/%s/%s' % (group, optType, name)

   def get(self):
      """Return all attributes of this config option.
      """
      return vsi.get(self._path)

   def set(self, value):
      """Set the config option to @value.
      """
      if self._strOpt and not isinstance(value, str):
         raise TypeError("%s: invalid value type (expected: 'str', got: %s)" %
                         (self._path, type(value)))

      if not self._strOpt and not isinstance(value, int):
         raise TypeError("%s: invalid value type (expected: 'int', got: %s)" %
                         (self._path, type(value)))

      return vsi.set(self._path, value)

   def reset(self):
      """Reset the config option to its default value.
      """
      return self.set(self.default)

   @property
   def description(self):
      """Return a description string for this config option.
      """
      return self.get()['description']

   @property
   def default(self):
      """Return thi config option's default value.
      """
      return self.get()['def']

   @property
   def current(self):
      """Return this config option's current value.
      """
      return self.get()['cur']

   @property
   def isHidden(self):
      """Return True if this config option is hidden.
      """
      return self._get()['isHidden'] == 1

   @property
   def min(self):
      """Return the minimum allowed value for this config option.
      """
      if self._strOpt:
         raise AttributeError("%s: string-type config option has no \"min\" "
                              "attribute" % self._path)
      return self.get()['min']

   @property
   def max(self):
      """Return the maximum allowed value for this config option.
      """
      if self._strOpt:
         raise AttributeError("%s: string-type config option has no \"max\" "
                              "attribute" % self._path)
      return self.get()['max']

   @property
   def valid(self):
      """Return the set of allowed characters for this config option.
      """
      if not self._strOpt:
         raise AttributeError("%s: integer-type config option has no \"valid\" "
                              "attribute" % self._path)
      return self.get()['valid']

def isUW32Enabled():
   """Return True if 32-bit UW are supported and enabled.
   """
   return VmkConfigOpt('User', 'UW32Enabled').current == 1

def vmkping(dst, netstack='defaultTcpipStack', count=3):
   """Ping utility for vmkernel TCP/IP stack
   """
   cmd = ['/sbin/vmkping', '++netstack={0}'.format(netstack),
          '-c', str(count), dst]
   p = subprocess.Popen(cmd,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        universal_newlines=True)
   out, err = p.communicate()
   res = p.returncode

   if res != 0:
      raise Exception(out)

   return out

def runCli(cmd, evalOutput=False):
   """Execute a localcli command.

   By default, this function returns the output of the excuted command formatted
   as a string that can be interpreted by python. Setting the evalOutput
   parameter to True will a evalutate this string and return a python object.

   This function raise an EsxcliError exception upon error.
   """
   if not hasattr(runCli, 'cliHandle'):
      runCli.cliHandle = esxclipy.EsxcliPy(True)

   status, out = runCli.cliHandle.Execute(cmd)
   if status != 0:
      raise EsxcliError(out)

   if evalOutput:
      try:
         out = eval(out)
      except Exception as e:
         raise EsxcliError("Failed to evaluate output of %s: %s" %
                           (' '.join(cmd), str(e)))
   return out

def getBuildInfo(key):
   """Parse .buildinfo to retrieve information about the current ESX build.
   """
   DOT_BUILDINFO = os.path.join(os.path.sep, 'etc', 'vmware', '.buildInfo')

   with open(DOT_BUILDINFO) as f:
      for line in f:
         var, value = tuple(line.strip('\n').split(':'))
         if var == key:
            return value

   return None

def getCln():
   """Get the current Perforce changeset number (CLN).
   """
   return int(getBuildInfo('CHANGE'))

def getBuildContext():
   """Get the current build context (ob or sb). Returns None if ESX was built
   locally.
   """
   vmtree = getBuildInfo('VMTREE')

   if re.search('sb-[0-9]+', vmtree) is not None:
      return 'sb'
   elif re.search('bora-[0-9]+', vmtree) is not None:
      return 'ob'
   return None

def getBuildType():
   """Get the current build type.
   """
   return getBuildInfo('VMBLD')

def getStagePath():
   """Get the current stage path.
   """
   return getBuildInfo('STAGEPATH')

def getSystemVersion():
   """Query localcli to retrieve ESX version information.
   """
   return runCli(['system', 'version', 'get'], evalOutput=True)

def getSystemUUID():
   """Get test system UUID.
   """
   return runCli(['system', 'uuid', 'get'], evalOutput=True)

def getHardwareUuid():
   """Query hardware platform UUID, which is from the SMBIOS table.
   """
   # VSI dmiInfo['uuid'] is an array of integers
   u = uuid.UUID(bytes=bytes(vsi.get('/hardware/bios/dmiInfo')['uuid']))
   return str(u).upper()

def getBuildNum():
   """Return ESX build number as a string.

   First try to get the build number from localcli. If it fails, fall back to
   parsing .buildinfo.
   """
   try:
      buildNum = getSystemVersion()['Build']
   except Exception:
      buildNum = None

   if buildNum is not None:
      m = re.search('\d+', buildNum)
      if m is not None:
         return m.group(0)

   buildNum = getBuildInfo('BUILDNUMBER')
   if buildNum is not None:
      return buildNum

   raise Exception("Failed to retrieve ESX build number!")

def getBuildId():
   """Return ESX build ID as a string.

   E.g. "00000" for developer builds, "sb-123456" for Sandbox builds, or
        "ob-457890" for Official builds.
   """
   DEV_BUILD_NUM = "00000"

   buildId = getBuildNum()
   if buildId != DEV_BUILD_NUM:
      builderType = getBuildContext()
      if builderType is not None:
         buildId = "%s-%s" % (builderType, buildId)

   return buildId

def getEsxVersion():
   """Return the current ESX version string.
   """
   version = getSystemVersion()['Version']
   return version.strip() if version is not None else None

def getCompatString(msg):
   """ Returns compatible string for python2 and python3
   """
   if sys.version_info[0] >= 3:
      # vsi expects unicode
      if not isinstance(msg, str):
         msg = msg.decode('utf-8')
   else:
      # vsi expects ascii
      msg = msg.encode('utf-8')

   return msg

def vmkLog(msg):
   """Log a message in the vmkernel logs.
   """
   msg = getCompatString(msg)
   vsi.set("/system/log", [msg])

def getFreeSpaceOnDataStore(volumeName):
   """Return the free space in MB given a volume name or
      -1 if no volume with this name can be found.
   """
   filesystems = runCli(["storage", "filesystem", "list"], evalOutput=True)

   for fs in filesystems:
      if fs['Volume Name'] == volumeName:
         return fs['Free'] / 1024 / 1024

   return -1

def getFolderSize(folder):
   """Return the space used by folder on disk. Result is returned in bytes.
      -1 if folder does not exist.
   """
   totalSize = 0

   if not os.path.isdir(folder):
      return -1

   for dirPath, _, fileNames in os.walk(folder):
      for f in fileNames:
         filePath = os.path.join(dirPath, f)
         totalSize += os.path.getsize(filePath)

   return totalSize

def getFsList(fsType=None, minFreeSpaceInMb=None, skipUnmounted=True,
              localOnly=False):
   """Return list of filesystems
   """
   filesystems = runCli(["storage", "filesystem", "list"], evalOutput=True)

   for fs in filesystems:
      if fsType is not None and not fs['Type'].startswith(fsType):
         continue
      if skipUnmounted and not fs['Mounted']:
         continue
      if minFreeSpaceInMb is not None:
         if fs['Free'] < (minFreeSpaceInMb * 1024 * 1024):
            continue
      yield fs

def getVmfsList(minFreeSpaceInMb=None, skipUnmounted=True, localOnly=False):
   """Iterate over VMFS filesystems.

   @skipUnmounted can be specified to ignore unmounted partitions.
   @localOnly can be set to ignore NFS-mounted partitions.
   """
   vmfsList = getFsList(fsType='VMFS-', minFreeSpaceInMb=minFreeSpaceInMb,
                        skipUnmounted=skipUnmounted)

   if localOnly:
      extents = runCli(['storage', 'vmfs', 'extent', 'list'], evalOutput=True)

   for vmfs in vmfsList:
      skip = False
      if localOnly:
         for extent in extents:
            if vmfs['Volume Name'] == extent['Volume Name']:
               devices = runCli(['storage', 'core', 'device', 'list', '-d',
                                 extent['Device Name']], evalOutput=True)
               assert len(devices) == 1
               device = devices[0]
               skip = not device['Is Local']
               break

      if not skip:
         yield vmfs

def getDefaultDatastore(minFreeSpaceInGB=None, localOnly=False):
   """Return a path to the first VMFS found that has enough free space.
   """
   if minFreeSpaceInGB is None:
      minFreeSpaceInMb = None
   else:
      minFreeSpaceInMb = minFreeSpaceInGB * 1024

   try:
      filesystems = getVmfsList(minFreeSpaceInMb=minFreeSpaceInMb,
                                localOnly=localOnly)
   except EsxcliError:
      return None

   try:
      return next(filesystems)['Mount Point']
   except StopIteration:
      return None

def getNumPcpus():
   """Return the number of PCPUs (threads) on the system.
   """
   cpuInfo = vsi.get('/hardware/cpu/cpuInfo')
   return cpuInfo['numCPUs']

def nfsMount(host, srcPath, name=None, readonly=True):
   """Mount an NFS share.
   """
   if name is None:
      name = "%s_%s" % (host, srcPath.replace('/', '.'))

   cmd = ['storage', 'nfs', 'add', '-H', host, '-s', srcPath, '-v', name]
   if readonly:
      cmd += ['-r']

   try:
      runCli(cmd)
   except Exception as e:
      match = re.search('already exported by a volume with the name'
                        '(?P<name>.+)', str(e), flags=re.IGNORECASE)
      if match is not None:
         volumeName = match.group('name')
         if volumeName == name:
            pass
         else:
            raise NfsAlreadyExportedError(volumeName)
      else:
         raise e

   nfsShare = os.path.join(os.sep, 'vmfs', 'volumes', name)
   assert os.path.isdir(nfsShare), "Failed to mount %s:%s to %s as %s is not a"\
          " directory" % (host, srcPath, name, nfsShare)
   return nfsShare

def nfsUnmount(volume):
   """Umount an NFS share.
   """
   volRoot = os.path.join(os.sep, 'vmfs', 'volumes')

   if volume.startswith(volRoot):
      volumeName = os.path.relpath(volume, volRoot)
   else:
      volumeName = volume

   runCli(['storage', 'nfs', 'remove', '--volume-name', volumeName])

def getHostIp(iface='vmk0'):
   """Return the host IPv4 address
   """
   ipInfo = runCli(['network', 'ip', 'interface', 'ipv4', 'get', '-i', iface],
                   evalOutput=True)
   ifaceInfo = ipInfo[0]
   assert ifaceInfo['Name'] == iface
   return ifaceInfo['IPv4 Address']

def getHostMemorySize():
   """Return host memory size in MB.
   """
   return runCli(['hardware', 'memory', 'get'],
                  evalOutput=True)['Physical Memory'] / (1024 * 1024)

def getSoftwareAcceptanceLevel():
   """Get the software acceptance level
   """
   try:
      level = runCli(['software', 'acceptance', 'get'], evalOutput=True)
   except EsxcliError as e:
      raise OSError("Failed to get the software acceptance level: %s" % e)
   return level

def setSoftwareAcceptanceLevel(level):
   """Set the software acceptance level
   """
   try:
      runCli(['software', 'acceptance', 'set', '--level=%s' % level])
   except EsxcliError as e:
      raise OSError("Failed to set the software acceptance level to %s: %s" %
                    (level, e))

def disableFireWall():
   """disable firewall
   """
   runCli('network firewall set --enabled=false'.split())

def enableFireWall():
   """enable firewall
   """
   runCli('network firewall set --enabled=true'.split())

def setFireWallRule(rule, value):
   """Set firewall rule to True/False
   """
   cli = 'network firewall ruleset set -r %s -e %s' % (rule, value)
   runCli(cli.split())

def disableFireWallRuleSSHClient():
   """disable firewall rule sshClient
   """
   setFireWallRule('sshClient', 'false')

def enableFireWallRuleSSHClient():
   """enable firewall rule sshClient. Keep it enabled to run in parallel.
   """
   setFireWallRule('sshClient', 'true')

def disableFireWallRuleHTTPClient():
   """disable firewall rule httpClient
   """
   setFireWallRule('httpClient', 'false')

def enableFireWallRuleHTTPClient():
   """enable firewall rule httpClient. Keep it enabled to run in parallel.
   """
   setFireWallRule('httpClient', 'true')

def checkFSType(path, fsType):
   """Check whether the given path lives on the given file system type.
   """
   if not os.path.exists(path):
      raise OSError('%s: no such file or directory.' % path)

   p = subprocess.Popen([VMKFSTOOLS, '-Ph', path], stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE)
   out, err = p.communicate()
   if p.returncode != 0:
      raise OSError('%s: Failed to get partition type (%s)' % (path, err))

   out = getCompatString(out)
   fsType = getCompatString(fsType)

   return fsType.lower() in out.split('\n')[0].lower()

def findVFAT(path=None):
   """If a volume was provided explicitly, then use that. Otherwise, return the
   /scratch partition.
   """
   if path is not None:
      if not checkFSType(path, 'vfat'):
         raise OSError('%s is not located on a vfat partition')
      return path

   try:
      scratch = os.readlink("/scratch")
   except Exception:
      raise OSError('/scratch partition not found')

   for fs in getFsList(fsType='vfat'):
      if fs['Mount Point'] == scratch:
         return scratch

   raise OSError("No VFAT found")

def createVmdk(vmdkPath, sizeInMb, diskFormat='thin'):
   """Create a new VMDK file with vmkfstools
   """
   cmd = [VMKFSTOOLS, '-c', '%dM' % sizeInMb, '-d', diskFormat, vmdkPath]

   p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
   out, err = p.communicate()
   if p.returncode != 0:
      raise OSError('Failed to create "%s", %s' % (vmdkPath, err))

def ramdiskAdd(name, mountPoint, minMem, maxMem, permissions, extraOpts=None):
   """Add visorfs ramdisk
   """
   os.makedirs(mountPoint, exist_ok=True)

   opts = ['visorfs', 'ramdisk', 'add', '-n', name, '-t', mountPoint,
           '-m', str(minMem), '-M', str(maxMem), '-p', permissions]
   cmd = ['system']
   if extraOpts is not None:
      # Use systemInternal to handle extra options
      cmd = ['--plugin-dir', '/usr/lib/vmware/esxcli/int/', 'systemInternal']
      cmd = cmd + opts + extraOpts
   else:
      cmd += opts
   runCli(cmd)

def ramdiskRemove(mountPoint):
   """Remove a visorfs ramdisk.
   """
   cmd = ['system', 'visorfs', 'ramdisk', 'remove', '-t', mountPoint]

   # The vmkernel Object Cache releases file objects asynchronously, which
   # requires to keep retrying until all internal ramdisk resources have
   # actually been reaped.
   for i in range(5):
      try:
         runCli(cmd)
         break
      except EsxcliError as e:
         if i < 4 and 'resource busy' in str(e):
            time.sleep(3)
         else:
            raise

def tardiskMount(tardisk, ramdisk):
   """Mount a tardisk onto a ramdisk
   """
   name = os.path.basename(tardisk)
   tardiskDir = '/tardisks' if ramdisk == "root" else '/tardisks.noauto'
   dest = os.path.join(tardiskDir, name)
   if os.path.exists(dest):
      raise OSError('%s: tardisk is already loaded' % name)

   shutil.move(tardisk, tardiskDir)

   if ramdisk != "root":
      vsi.set('/system/visorfs/tardisks/%s/attachToRamdisk' % name, ramdisk)

def tardiskUnmount(name, ramdisk):
   """Unmount a tardisk
   """
   tardiskDir = '/tardisks' if ramdisk == "root" else '/tardisks.noauto'
   os.unlink(os.path.join(tardiskDir, name))

def installVibs(vibList=None, options=""):
   """Install VIB package on ESX
   """
   try:
      for vib in vibList:
         #if not os.path.exists(vib):
         #   raise OSError('Failed to find vib %s' % vib)
         runCli(['software', 'vib', 'install', '-v', vib, options])
   except Exception as e:
      raise Exception('Failed to install vib %s. Got error %s' % (vib, e))

def removeVibs(vibList=None):
   """Remove VIB packages on ESX.
   """
   try:
      for vib in vibList:
         runCli(['software', 'vib', 'remove', '-n', vib])
   except Exception as e:
      raise Exception('Failed to remove vib %s. Got error %s' % (vib, e))

def getVmkBootOptions():
   """Return all vmkernel boot options as a dictionary.
   """
   argv = vsi.get('/system/bootCmdLine')['bootCmdLineStr'].split()

   # argv[0] is the path to vmkBoot. Should bot be reported as a boot option.
   argv.pop(0)

   opts = {}
   for opt in argv:
      key, _, val = opt.partition('=')
      if val == "":
          val = None
      opts[key] = val
   return opts

def isEsxInstaller():
   """Return True if ESX is running in Installer mode.
   """
   return 'runweasel' in getVmkBootOptions()

def isScratchPartPersistent():
   """Checks if the scratch partition is persistent.

   If scratch is persistent, then /scratch points to /vmfs/volumes/...

   TODO: This function needs to be moved to the systemStorage library as part
         of the ESX SystemStorage project -- rvoltz.
   """
   try:
      path = os.path.realpath(SCRATCH_PARTITION)
      return os.path.exists(path) and path.startswith('/vmfs/volumes/')
   except OSError:
      return False

def getFQDN():
   """Return ESX Fully Qualified Domain Name (FQDN).
   """
   hostNames = runCli(["system", "hostname", "get"], True)
   return hostNames['Fully Qualified Domain Name']

