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

import os
import platform
import sys

MIB = 1024 * 1024
# Used for sys.path manipulation.
PYTHON_VER_STR = 'python%u.%u' % (sys.version_info.major,
                                  sys.version_info.minor)

_JSON_SCHEMA_DIR = 'jsonschemadir'
_CERTS_DIRS = 'certsdirs'
_SCHEMA_DIR = 'schemadir'

ESX_BASE_VIB = 'esx-base'
ESXIO_BASE_VIB = 'esxio-base'
BASE_VIBS = (ESX_BASE_VIB, ESXIO_BASE_VIB)

# The FSS values below do not mean anything during builds i.e. it will be False
# always even if the actual state is enabled.
# For EPK, FSS works only if the EPK build is from main branch.
# To change the state of an FSS switch for an EPK installation,
# use feature-state-util tool. This will modify the value of the switch present
# in /opt/vmware/esxpackagingkit/vsphereFeatures/vsphereFeatures.cfg, eg. if
# we want to enable a feature, change the corresponding "disabled" state to
# "enabled". In order to use feature-state-util, VMWARE_CFG_DIR must be set.
# Set it to point to epk base directory.

fssNames = ['ESXLiveUpdate', 'ESXLiveUpdateStage1', 'ESXZstdCompression',
            'InfraPartialMMode', 'LiveUpdate',
            'PersonalityManagerComponentRemoval',
            'PersonalityManagerDepotRecall']
fssVals = []
try:
   import featureState
   featureState.init(False)
   for fss in fssNames:
      try:
         fssVals.append(getattr(featureState, fss))
      except AttributeError:
         # Legacy ESXi.
         fssVals.append(False)
except ImportError:
   fssVals = [False] * len(fssNames)

(QUICKPATCH_STAGE2_ENABLED, QUICKPATCH_STAGE1_ENABLED, ZSTD_COMPRESSION_ENABLED,
 INFRAPARTIALMMODE_ENABLED, LIVEUPDATE_ENABLED,
 PERSONALITY_MANAGER_COMPONENT_REMOVAL_ENABLED,
 PERSONALITY_MANAGER_DEPOT_RECALL_ENABLED) = fssVals

# Enabled Features
PERSONALITY_MANAGER_COMPONENT_REMOVAL_ENABLED = True
# Bug 3394608: The FSSes have been enabled, but must be hard-coded to True for
# build time consumption until the removal.
INFRAPARTIALMMODE_ENABLED = True
QUICKPATCH_STAGE1_ENABLED = True
QUICKPATCH_STAGE2_ENABLED = True

NON_UNIFIED_IMAGE_TARGETS = set(['developer', 'esxallcomm', 'esxall-crypto2',
   'esxall-gccnext', 'esxall-hostd-malloc-bt', 'esxall-openssl3',
   'esxall-tools-compat', 'esxallasan', 'esxallasanuw', 'esxallcomm',
   'esxallcov', 'esxallcov-agents', 'esxallcov-hostd', 'esxallcov-settingsd',
   'esxallcov-ulm', 'esxallcov-vmk', 'esxallcov-vmkmod', 'esxallcov-vmk-hostd',
   'esxallcov-vmm', 'esxallcov-vmx', 'esxallcov-vsan', 'esxallcov-vvold',
   'esxallsymdb', 'esxarm64', 'esxarm64asan', 'esxarm64asan-vhe', 'esxarm64cov',
   'esxarm64cov-ulm', 'esxarm64cov-vmk', 'esxarm64symdb', 'esxarm64-openssl3',
   'esxarm64-vhe', 'esxcore', 'esxcorearm64', 'esxcoreriscv64', 'esxio',
   'esxioasan', 'esxiocov', 'esxiocov-vmk', 'esxiocov-hostd',
   'esxiocov-vmk-hostd', 'esxiox86', 'esxiox86cov', 'esxiox86cov-hostd',
   'esxiox86cov-vmk-hostd', 'esxio-vhe', 'serverarm64', 'visorpxe'])

# These payloads encapsulate state, they can change and hence their
# size and checksum may not match the one specified in descriptor.
CHECKSUM_EXEMPT_PAYLOADS = ["useropts", "features", "jumpstrt"]

# The VIB payloads for which xz is used during compression. Payload names
# are truncated at 8 characters.
# Payload           : Manifest Name in bora/install/vmvisor/packaging
# s                 : sys
# sb                : sys-boot
# vmx               : vmx
# vim               : vim
# clusters          : clusterstore
# trx               : trx
# vdfs              : vdfs
# vsan              : vsan
# vsanmgmt          : vsan-health
# infravis          : infravisor
# drivervm          : drivervm-gpu-base
XZ_PAYLOADS = ("s.v00", "sb.v00", "vmx.v00", "vim.v00",
               "clusters.v00", "crx.v00", "trx.v00",
               "vdfs.v00", "vsan.v00", "vsanmgmt.v00", "infravis.v00",
               "drivervm.v00")

def isEsxioAndIsEsxioX86():
   """Returns two booleans: whether the system is ESXio and whether it is ESXio
      on X86.
   """
   try:
      import vmkdefs
      return vmkdefs.vmx86_esxio, vmkdefs.vmx86_esxio and vmkdefs.vm_x86_64
   except (ImportError, AttributeError):
      # Legacy ESXi or vCenter/Linux.
      return False, False

IS_ESXIO, IS_ESXIO_X86 = isEsxioAndIsEsxioX86()
IS_ESX_ESXIO = getattr(platform.uname(), 'system', None) == 'VMkernel'
IS_ESXCLI = (IS_ESX_ESXIO and os.path.basename(sys.argv[0]) in
   ('esxcli-software', 'esxcli-softwareinternal'))

# Name of the patcher component in ESXi.
PATCHER_COMP_NAME = 'esxio-update' if IS_ESXIO else 'esx-update'

def _reImportSystemStorage():
   """Re-import all systemStorage and dependent modules required for an upgrade.
      This assumes sys.path contains the proper systemStorage.zip path.
   """
   # If a module has dependency on another module in the list, it must comes
   # later.
   # - borautils, entropy, esxutils and uefi use only Python lib and core infra
   #   such as vsi.
   # - advcfg, coredump and vmSyslogUtils import esxutils.
   # - systemStorage imports all 4 others.
   # - lifecycle.utils imports borautils.slock, if not re-imported,
   #   the locking retry mechanism will fail.
   MODULE_ORDER = ('borautils', 'entropy', 'esxutils', 'uefi', 'advcfg',
                   'coredump', 'vmSyslogUtils', 'systemStorage',
                   'lifecycle.utils')

   import importlib
   for name in MODULE_ORDER:
      found = False
      for n, m in sorted(sys.modules.items(), key=lambda x: x[0]):
         # Must use sorted list to always reload the root first.
         if n.startswith(name):
            importlib.reload(m)
            found = True
         elif found:
            # Stop as all matches have been reloaded.
            break

# Whether esximage is loaded as part of Patch the Patcher in esxcli/vLCM/VUM,
# incl. VUM ISO upgrade where the new esximage.zip is loaded.
IS_PATCH_THE_PATCHER = False

def _configurePatcher():
   """Configure esximage lib within a patcher mounted by esxcli/vLCM, or within
      esximage.zip during VUM ISO upgrade.
   """
   global IS_PATCH_THE_PATCHER

   modulePath = os.path.dirname(os.path.abspath(__file__))
   patcherPrefix = '/tmp/%s-' % PATCHER_COMP_NAME

   if patcherPrefix in modulePath and not '.zip' in modulePath:
      # esxcli/vLCM upgrade mounts new esx-update VIB in /tmp, in this case
      # esximage path should be like:
      # /tmp/esx-update-<pid>/lib64/python<ver>/site-packages/vmware/esximage
      IS_PATCH_THE_PATCHER = True

      # Amend sys.path to include paths for importing other new libraries:
      # - uefi module for UEFI boot option manipulation.
      # - systemStorage.upgradeUtils/esxboot and their dependencies for boot
      #   disk partition layout manipulation and bootloader installation.
      # - esxutils for misc tools
      sitePkgPath = os.path.normpath(os.path.join(modulePath, '..', '..'))
      mountRoot = os.path.normpath(os.path.join(sitePkgPath, '..', '..', '..'))

      sysStorageZipPath = os.path.join(mountRoot, 'usr', 'lib', 'vmware',
                                       'esxupdate', 'systemStorage.zip')

      for path in (sitePkgPath, sysStorageZipPath):
         if not path in sys.path:
            # The older esximage lib may have added some of these paths.
            sys.path.insert(0, path)

      # Legacy esximage could have already imported the local systemStorage
      # from /lib64, we cannot rely on it since the lib can change.
      # Instead, re-import the new version from systemStorage.zip mounted
      # in /tmp.
      _reImportSystemStorage()

      # Set the new XML schema and CA store paths. When tardisks are mounted
      # VIB signature and tardisk checksums are verified, this means we can
      # trust the CA store shipped within them.
      usrSharePath = os.path.join(mountRoot, 'usr', 'share')
      params = {
         _CERTS_DIRS: [
            # Trust VIBs that can be verified by either current or new CA store.
            # This allows internal upgrades from an image with test-cert to one
            # with mixed official/test signings.
            os.path.join(usrSharePath, 'certs'),
            os.path.join(os.path.sep, 'usr', 'share', 'certs'),
         ],
         _SCHEMA_DIR: os.path.join(usrSharePath, 'esximage', 'schemas'),
      }
      Configure(**params)
   elif '.zip' in modulePath and 'vuaScript' in modulePath:
      # During ISO upgrade, VUA folder contains esximage.zip which includes
      # XML schemas, extract them to configure schema dir. CA store path will
      # remain at its current path.
      from zipfile import is_zipfile, ZipFile

      IS_PATCH_THE_PATCHER = True

      zipPath = modulePath
      while not (zipPath.endswith('.zip') or zipPath == os.path.sep):
         zipPath = os.path.dirname(zipPath)

      if not is_zipfile(zipPath):
         # Unexpected zip path or zip not exist.
         return

      workDir = os.path.dirname(zipPath)
      try:
         schemaPrefix = os.path.join('usr', 'share', 'esximage', 'schemas')
         with ZipFile(zipPath, 'r') as z:
            for i in z.infolist():
               if schemaPrefix in i.filename:
                  z.extract(i, workDir)

         params = {
            _SCHEMA_DIR: os.path.join(workDir, schemaPrefix),
         }
         Configure(**params)
      except Exception:
         # Do not panic, most likely current schema will just work.
         pass

def GetEsxImageUserVars():
   """Get the EsxImage UserVars to be used with Configure().
   """
   # ESX-only import
   from vmware import runcommand
   ADVCFG = '/sbin/esxcfg-advcfg'

   opts = dict()
   for userVar, key in (("EsximageNetTimeout", "nettimeout"),
                        ("EsximageNetRetries", "netretries"),
                        ("EsximageNetRateLimit", "netratelimit")):
      try:
         res, out = runcommand.runcommand(
                        [ADVCFG, '-q', '-g', '/UserVars/' + userVar])
         if res == 0 and out:
            opts[key] = int(out.strip())
      except Exception:
         opts[key] = None
   return opts

def Configure(**kwargs):
   """This function is used to configure various aspects of the module's
      operation. The following keyword arguments are accepted:
         * nettimeout    - A positive integer or float giving the amount of time
                           to wait for reads from a connection to an HTTP, HTTPS
                           or FTP server. May also be None or 0, which disables
                           the timeout.
         * netretries    - A positive integer specifying the number of times to
                           retry a connection to an HTTP, HTTPS or FTP server.
                           A value of 0 causes infinite retries. This may also
                           be None, which disables retrying.
         * netratelimit  - A positive integer specifying, in bytes per second,
                           the maximum bandwidth to use for HTTP, HTTPS and FTP
                           downloads.
         * certsdir      - Specifies a path to a directory containing the
                           certificates to be used for acceptance level
                           verification.
         * schemadir     - Specifies a path to a directory containing the
                           schemas to be used for acceptance level verification
                           and schema validation.
         * jsonschemadir - Specifies a path to a directory containing the
                           json schemas to be used for schema validation.
   """
   def checkDirArg(dirArg, argName):
      if not isinstance(dirArg, str) and not isinstance(dirArg, bytes):
         raise ValueError("'%s' input must be a string" % argName)
      if not os.path.isdir(dirArg):
         raise ValueError("'%s' is not a directory or does not exist"
                          % dirArg)

   if "nettimeout" in kwargs:
      from . import Downloader
      Downloader.SetTimeout(kwargs.pop("nettimeout"))
   if "netretries" in kwargs:
      from . import Downloader
      Downloader.SetRetry(kwargs.pop("netretries"))
   if "netratelimit" in kwargs:
      from . import Downloader
      Downloader.SetRateLimit(kwargs.pop("netratelimit"))
   if _JSON_SCHEMA_DIR in kwargs:
      from .Utils import JsonSchema
      schemaDir = kwargs.pop(_JSON_SCHEMA_DIR)
      checkDirArg(schemaDir, _JSON_SCHEMA_DIR)
      JsonSchema.SCHEMA_ROOT = schemaDir
   if _SCHEMA_DIR in kwargs:
      from . import Bulletin, ImageProfile, Vib
      schemaDir = kwargs[_SCHEMA_DIR]
      checkDirArg(schemaDir, _SCHEMA_DIR)
      for module in (Bulletin, ImageProfile, Vib):
         module.SCHEMADIR = schemaDir

   al_args = dict()
   for key in (_CERTS_DIRS, _SCHEMA_DIR):
      if key in kwargs:
         al_args[key] = kwargs.pop(key)

   if al_args:
      from . import AcceptanceLevels
      AcceptanceLevels.Initialize(**al_args)

   if kwargs:
      raise TypeError("configure() got unexpected keyword argument(s): %s"
                      % ", ".join(kwargs))

def hasDpuVapi():
   try:
      from com.vmware.esx.settings_daemon_client import \
         DataProcessingUnitsCompliance
      return True
   except ImportError:
      return False

def isHigherEsx8000():
   """Returns whether ESXi version is higher than 8.0.0-0.x.y.
   """
   try:
      from vmware import vsi
      verInfo = vsi.get('/system/version')
      verStr = verInfo['releaseVersionStr']
      return int(verStr.split('.')[0]) >= 8 and not verStr.startswith('8.0.0-0')
   except (ImportError, KeyError):
      # Legacy ESXi or vCenter/Linux.
      return False

_BUILD_INFO = '/etc/vmware/.buildInfo'

def isInternalEpk():
   """Returns whether esximage is running in the internal EPK build.
   """
   # Rule out ESXi/VC.
   if not os.path.isfile(_BUILD_INFO):
      # esximage module path is: site-packages\esximage.zip\vmware\esximage\
      modulePath = os.path.dirname(os.path.abspath(__file__))
      if 'esximage.zip' in modulePath:
         epkRootPath = os.path.dirname(os.path.dirname(os.path.dirname(
            os.path.dirname(modulePath))))
         if os.path.exists(os.path.join(epkRootPath, 'bin', 'depotAuthor')):
            # depotAuthor only exists for non-RPM EPK builds
            return True
   return False

IS_INTERNAL_EPK = isInternalEpk()

def isUnifiedBuildTarget():
   """Returns whether esximage is running in a build of an unified build target.
   """
   try:
      with open(_BUILD_INFO, 'r') as buildInfo:
         for line in buildInfo:
            key, val = line.strip().partition(':')[::2]
            if key == 'GOBUILDTARGET':
               return val not in NON_UNIFIED_IMAGE_TARGETS
   except Exception:
      # Default to True.
      return True

IS_UNIFIED_BUILD_TARGET = isUnifiedBuildTarget()

HAS_DPU = True
if IS_ESX_ESXIO:
   def containsDPU():
      """Returns True if there is a managed DPU present.
      """
      from .ESXioImage.DpuLib import getManagedDpuInfo
      try:
         dpuList = getManagedDpuInfo()
         return len(dpuList) > 0
      except Exception:
         # There might be an issue with DPU management. Marking DPU presence
         # False to avoid execution of DPU ops.
         return False
   HAS_DPU = containsDPU()
   _configurePatcher()

# ESXi 8.0.0-0 has DPU VAPI but should not perform DPU operation.
ALLOW_DPU_OPERATION = (not IS_ESXIO and HAS_DPU and hasDpuVapi() and isHigherEsx8000())
