# Copyright (c) 2021 VMware, Inc. All rights reserved.
# VMware Confidential

import os
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'

fssNames = ['ConfigStoreStage2', 'LiveUpdate']
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)

(CONFIG_STORE_ENABLED, LIVEUPDATE_ENABLED) = fssVals

SYSTEM_STORAGE_ENABLED = True

# 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"]

def _reImportSystemStorage():
   """Re-import all systemStorage modules assuming sys.path contains the proper
      systemStorage.zip path.
   """
   import importlib

   # importlib does not recognize "from" imports without explictly importing
   # the parent module once.
   import systemStorage
   importlib.reload(systemStorage)

   # Re-import each sub-module
   subModules = [module for name, module in sys.modules.items()
                 if 'systemStorage.' in name]
   for module in subModules:
      importlib.reload(module)

def _configurePatcher():
   """Configure esximage lib within a patcher mounted by esxcli/vLCM, or within
      esximage.zip during VUM ISO upgrade.
   """
   modulePath = os.path.dirname(os.path.abspath(__file__))

   if '/tmp/' 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

      # 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

      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))


_configurePatcher()
