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

"""Interfaces that call vAPI or esxcli on all DPUs managed by this host.
"""

from datetime import datetime
import itertools
import json
import logging

from com.vmware.esx.settings_daemon_client import (
   DataProcessingUnitsCompliance, DataProcessingUnitInfo,
   DataProcessingUnitCompliance)
from com.vmware.esx.settings_daemon import data_processing_unit_client as dpuc
from com.vmware.esx.software_client import InstalledComponents

from .DpuLib import getDpuAlias, getDpuID, getDpuInfoFromId, getManagedDpuInfo
from .DpuTask import (ApplySolutionTask, ApplyWorkflow, ComponentApplyTask,
                      VapiDpuListTask, DpuVapiEsxcliNonTask, DpuEsxcliTask,
                      RemoveSolutionTask, ScanTask)
from .TaskBase import RunnableGroup

from .Constants import *
from .Utils import createNotification, createDummyObjectForType

from .. import IS_ESXIO
from ..ImageManager import HostDepot
from ..Utils.Misc import LogLargeBuffer

from ..Errors import DpuInfoError, DpuOperationError, NormalExit

log = logging.getLogger(__name__)

# Expand list to a repeating series of command-line options, opt must have
# space on both sides.
_optionJoin = lambda opt, vals: opt + opt.join(vals)

USE_ESXCLI_VAPI = hasattr(dpuc, 'Esxcli')

class ImageOperation(object):
   """ Base class for image operations.
   """
   # Whether the operation has void result.
   HAS_VOID_RESULTS = True

   # Whether the operation depends on depots.
   USE_DEPOTS = True

   def __init__(self, depots, parentTask=None, extraVibs=None):
      """ Constructor for image operation.

          depots: The depots to support this image operation.
          parentTask: The parent task for ESXi host scan.
          extraVibs: Standalone VIB URLs/paths, to be included into the
                     host depot.
      """
      self.APPLY_CLASSES = (ApplyOperation, ComponentApplyOperation,
                            ApplySolutionOperation, RemoveSolutionOperation,
                            ComponentApplyEsxcliOperation,
                            ComponentRemoveOperation,
                            ProfileInstallUpdateOperation,
                            ApplyEsxcliOperation, VibInstallUpdateOperation,
                            VibRemoveOperation, AcceptanceGetSetOperation)

      # The "get/list" operations need not be blocked if one or more DPUs are
      # unavailable. This is useful during collection of log bundle as well.
      self._dpus = getManagedDpuInfo(skipUnavailableDpu=type(self) not in
                                     self.APPLY_CLASSES)
      self._depots = depots
      self._extraVibs = extraVibs
      self._hostDepot = None
      self._localDepots = []
      self._parentTask = parentTask
      self._dpuResults = self.initDpuResults()
      self._dpuTasks = []
      self._componentsToDownload = None

   def _getDpuTasks(self, taskCls, preUrlArgs, postUrlArgs):
      """Returns a list of tasks represented by taskCls to be run on the DPU.

         Parameters:
            * taskCls     - DPU task class whose object will be returned.
            * preUrlArgs  - List of arguments to the taskCls constructor that
                            appear before the depot url parameter.
            * postUrlArgs - List of arguments to the taskCls constructor that
                            appear after the depot url parameter.
      """
      dpuTasks = []
      for dpu, depotUrl in zip(self._dpus, self._localDepots):
         if postUrlArgs:
            task = taskCls(dpu, *preUrlArgs, [depotUrl], *postUrlArgs)
         else:
            task = taskCls(dpu, *preUrlArgs, [depotUrl])
         dpuTasks.append(task)
      return dpuTasks

   def checkNotApplicable(self):
      """ Don't execute image operation for:
             ESXio device
             no DPU exists
             when using host depot, neither depots nor extraVibs is filled
      """
      return (IS_ESXIO or not self._dpus or
              (self.USE_DEPOTS and not self._depots and not self._extraVibs))

   def createHostDepot(self):
      """ Create the local host depot.
      """
      self._hostDepot = HostDepot.HostDepot()
      self._hostDepot.createHostDepotFromList(self._depots,
            compSpecs=self._componentsToDownload,
            extraVibs=self._extraVibs)
      self._localDepots = HostDepot.getHostDepotURLs(self._dpus)
      return bool(self._localDepots)

   def run(self):
      try:
         if self.checkNotApplicable():
            return True if self.HAS_VOID_RESULTS else None

         if self.USE_DEPOTS and not self.createHostDepot():
            log.error('Create host depot for DPU image operation failed.')
            return False if self.HAS_VOID_RESULTS else None

         if self._parentTask:
            args = [str(len(self._dpus))]
            notif = createNotification(DPUOperationStart, DPUOperationStart,
                                       args, args)
            self._parentTask.updateNotifications([notif])

         self._dpuTasks = self.createDpuTasks()
         runnableGroup = RunnableGroup(self._dpuTasks)
         runnableGroup.run()
         self.finalizeOperation()
         return (runnableGroup.succeeded() if self.HAS_VOID_RESULTS
                                           else self._dpuResults)
      except Exception as e:
         log.error('Cannot run image operation on dpus: %s', str(e))
         return False if self.HAS_VOID_RESULTS else None

   def initDpuResults(self):
      """ API to create and intialize the DPU results.
      """
      pass

   def finalizeOperation(self):
      """API to finalize the operation: post process DPU results.
      """
      pass

class ComponentApplyOperation(ImageOperation):
   """ Component apply on all DPUs.
   """

   def __init__(self, components, depots):
      """ Constructor for component apply operation.

          components: The components to be applied, each entry must be of format
                      name:version.
          depots: The depots to support this image operation.
      """
      super(ComponentApplyOperation, self).__init__(depots)

      # Only support name:version in VAPI route, esxcli that supports name-only
      # input will run via ssh.
      componentsDict = dict()
      for comp in components:
         try:
            name, version = comp.split(':')
            componentsDict[name] = version
         except ValueError:
            log.warning('Ignoring component %s without version for DPU', comp)
      self._components = componentsDict
      self._componentsToDownload = list(componentsDict.keys())

   def checkNotApplicable(self):
      """ Check not applicable cases.
      """
      return super(ComponentApplyOperation, self).checkNotApplicable() or \
             not self._components

   def createDpuTasks(self):
      """ Create DpuTasks for each DPU.
      """
      preUrlArgs = [self._components]
      postUrlArgs = None
      return self._getDpuTasks(ComponentApplyTask, preUrlArgs, postUrlArgs)

class ScanOperation(ImageOperation):
   HAS_VOID_RESULTS = False
   PTP_COMPONENTS = ['esxio-update']

   """ Scan image on all DPUs.
   """
   def __init__(self, softwareSpec, depots, relatedComponents, parentTask=None):
      """ Constructor for image scan operation.

          softwareSpec: The desired image.
          depots: The depots to support this image operation.
          relatedComponents: The (name, version) pairs of components related
                             to the desired image.
          parentTask: The parent task for ESXi host scan.
      """
      super(ScanOperation, self).__init__(depots, parentTask)
      self._softwareSpec = softwareSpec
      self._componentsToDownload = [(n, v) for n, v in relatedComponents
                                    if n in self.PTP_COMPONENTS]

   def initDpuResults(self):
      """ Create and intialize the DPU results.
      """
      dpuResults = createDummyObjectForType(DataProcessingUnitsCompliance)
      dpuInfos = dpuResults.data_processing_unit_info
      for dpu in self._dpus:
         dpuId = getDpuID(dpu)
         dpuInfo = createDummyObjectForType(DataProcessingUnitInfo)
         dpuInfo.alias = dpuId
         dpuInfos[dpuId] = dpuInfo

         # Init all DPU compliance as UNAVAILABLE.
         compliance = createDummyObjectForType(DataProcessingUnitCompliance)
         compliance.scan_time = datetime.utcnow()
         compliance.status = UNAVAILABLE
         compliance.impact = ComplianceImpact.UNKNOWN
         compliance.base_image.status = UNAVAILABLE
         compliance.add_on.status = UNAVAILABLE
         dpuResults.compliance[dpuId] = compliance
      return dpuResults

   def createDpuTasks(self):
      """ Create scan task for each DPU.
      """
      preUrlArgs = [self._softwareSpec]
      postUrlArgs = [self._parentTask, self._dpuResults]
      return self._getDpuTasks(ScanTask, preUrlArgs, postUrlArgs)

   def finalizeOperation(self):
      """API to finalize the operation. Categorize the DPUs based on the
         compliance status.
      """
      dpuCompliances = self._dpuResults.compliance
      for dpuId, res in dpuCompliances.items():
         if res.status == COMPLIANT:
            self._dpuResults.compliant_data_processing_units.add(dpuId)
         elif res.status == NON_COMPLIANT:
            self._dpuResults.non_compliant_data_processing_units.add(dpuId)
         elif res.status == INCOMPATIBLE:
            self._dpuResults.incompatible_data_processing_units.add(dpuId)
         elif res.status == UNAVAILABLE:
            self._dpuResults.unavailable_data_processing_units.add(dpuId)

class ApplyOperation(ImageOperation):
   """ Apply desired image to all DPUs.
   """

   def __init__(self, softwareSpec, depots, componentsToDownload,
                parentTask=None):
      """ Constructor.

          softwareSpec: The desired image to be applied.
          depots: A depot list.
          componentsToDownload: Components related to the desired image.
          parentTask: The parent task of ESXi host apply.
      """
      super(ApplyOperation, self).__init__(depots, parentTask)
      self._softwareSpec = softwareSpec
      self._componentsToDownload = componentsToDownload

   def createDpuTasks(self):
      """ Create ApplyWorkflow for each DPU.
      """
      preUrlArgs = [self._softwareSpec]
      postUrlArgs = [self._parentTask]
      return self._getDpuTasks(ApplyWorkflow, preUrlArgs, postUrlArgs)

class ApplySolutionOperation(ImageOperation):
   """ Apply solution on all DPUs.
   """

   def __init__(self, solutions, depots, parentTask=None):
      """ Constructor.

          solutions: The desired solutions to be applied.
          depots: A depot list.
          parentTask: The parent task of ESXi host solution apply.
      """
      super(ApplySolutionOperation, self).__init__(depots, parentTask)
      self._solutions = solutions
      self._componentsToDownload = []
      for solComps in solutions.values():
         for comp in solComps['components']:
            self._componentsToDownload.append(comp['component'])

   def createDpuTasks(self):
      """ Create ApplySolutionTask for each DPU.
      """
      preUrlArgs = [self._solutions]
      postUrlArgs = [self._parentTask]
      return self._getDpuTasks(ApplySolutionTask, preUrlArgs, postUrlArgs)

def applySolutionOnDPUs(solutions, depots, parentTask):
   """ Apply solution on DPUs.

      solutions: The solution section in a desired image spec.
      depots: A depot file/URL list.
      parentTask: The parent task for ESXi host solution apply.

      Return: NormalExit on success; otherwise, DpuOperationError.
   """
   try:
      applySolOp = ApplySolutionOperation(solutions, depots, parentTask)
      if applySolOp.run():
         return NormalExit()
      else:
         return DpuOperationError('Failed to apply solution.')
   except DpuInfoError as e:
      log.exception('Failed to get DPU info: %s', e)
      if parentTask:
         notif = createNotification(DPUInfoError, DPUInfoError, [], [], ERROR)
         parentTask.updateNotifications([notif])
      return DpuOperationError('Failed to apply solution on DPUs: %s' % str(e))
   except Exception as e:
      msg = 'Failed to apply solution on DPUs: %s' % str(e)
      log.exception(msg)
      return DpuOperationError(msg)

class RemoveSolutionOperation(ImageOperation):
   """ Remove solutions on all DPUs.
   """

   USE_DEPOTS = False

   def __init__(self, solutionNames, parentTask=None):
      """ Constructor.

          solutionNames: The solutions to be removed.
          parentTask: The parent task of ESXi host solution removal.
      """
      super(RemoveSolutionOperation, self).__init__(None, parentTask)
      self._solutionNames = solutionNames


   def createDpuTasks(self):
      """ Create RemoveSolutionTask for each DPU.
      """
      return [RemoveSolutionTask(dpu, self._solutionNames, self._parentTask)
              for dpu in self._dpus]

class VapiListOperation(ImageOperation):
   """ Provide list of all the component or vib available on DPU(s)
   """
   USE_DEPOTS = False

   def __init__(self, isVibList, orphanVibs=False):
      """ Constructor.

          vibList: Flag to create compoent or vib list operation.
          orphanVibs: Flag to get orphan vibs list from Dpu.
      """
      super(VapiListOperation, self).__init__(None)
      self._isVibList = isVibList
      self._orphanVibs = orphanVibs

   def createDpuTasks(self):
      """ Create task for each DPU.
      """
      return [VapiDpuListTask(dpu, self._dpuResults, self._isVibList,
                              self._orphanVibs) for dpu in self._dpus]

   def initDpuResults(self):
      """ Create and intialize the DPU results.
      """
      dpuOutputList = {}
      for dpu in self._dpus:
         dpuId = getDpuID(dpu)
         dpuOut = createDummyObjectForType(
            InstalledComponents.InstalledComponentInfo)
         dpuOutputList[dpuId] = dpuOut
      return dpuOutputList

   def finalizeOperation(self):
      """API to finalize the operation.
      """
      if self._dpuResults:
         dpuListiOutput = []
         for dpuId, dpuRes in self._dpuResults.items():
            dpuInfo = getDpuInfoFromId(self._dpus, dpuId)
            dpuAlias = getDpuAlias(dpuInfo)
            if self._isVibList:
               for vibinfo in dpuRes:
                  _dpuResult = {
                     'vib' : vibinfo.get_field('vib'),
                     'name' : vibinfo.get_field('name'),
                     'version' : vibinfo.get_field('version'),
                     'platforms': {dpuAlias}
                  }
                  dpuListiOutput.append(_dpuResult)
            else:
               for component, compInfo in dpuRes.items():
                  _dpuResult = {
                     'component' : component,
                     'display_name': compInfo.get_field('display_name'),
                     'version' : compInfo.get_field('version'),
                     'display_version': compInfo.get_field('display_version'),
                     'platforms': {dpuAlias}
                  }
                  dpuListiOutput.append(_dpuResult)
         self._dpuResults = dpuListiOutput

class EsxcliOperation(ImageOperation):
   """ Run esxcli command on all DPUs.
   """
   TIMEOUT = 100

   def __init__(self, command, isLocalcli=False, isInternal=False, depots=None,
                extraVibs=None, extraArgs=None):
      """ Constructor.

          Parameters:

          command: A subcommand of esxcli software command without options.
          isLocalcli: set to True when localcli should be used instead of
                      esxcli.
          isInternal: set to True when softwareinternal namespace should be
                      used, isLocalcli must be set together when set.
          depots: A depot file/URL list, each will be expanded to -d <depot>
                  option and appended to the final command.
          extraVibs: Standalone VIB URLs/paths, to be included into the host
                     depot.
          extraArgs: A key-value dict of args, value can be empty indicating
                     a flag, e.g. {'-p': '<profileName>', '-f': None}; each
                     arg will be appended to the final command.
      """
      super(EsxcliOperation, self).__init__(depots, extraVibs=extraVibs)
      self._command = command
      self._commandStr = command
      self._results = dict()
      self._isLocalcli = isLocalcli
      self._isInternal = isInternal
      self._extraArgs = extraArgs
      if extraArgs:
         for option, val in extraArgs.items():
            self._command += ' %s "%s"' % (option, val) if val else ' %s' % option

   def _addExtraCommandArgs(self, dpu=None, depotUrl=None):
      """Add depot argument to the command before task creation.
         This function gets called from the sub-class in ssh path.
         dpu argument is unused here.
      """
      if self.USE_DEPOTS:
         if not USE_ESXCLI_VAPI:
            # Add "-d <URL>" option to command now host depot has been created.
            depotOpt = ' -d %s' % depotUrl
            if depotOpt not in self._command:
               self._command += depotOpt

   def createDpuTasks(self):
      """ Create DPU task for each DPU.
      """
      if USE_ESXCLI_VAPI:
         dpuVapiEsxcliNonTaskList = []
         for dpu, depotUrl in itertools.zip_longest(self._dpus,
                                                    self._localDepots,
                                                    fillvalue=None):
            self._addExtraCommandArgs(dpu=dpu)
            dpuVapiEsxcliNonTask = DpuVapiEsxcliNonTask(dpu, self._commandStr,
                                       self._isLocalcli, self._isInternal,
                                       [depotUrl] if self.USE_DEPOTS else None,
                                       self._extraArgs, timeout=self.TIMEOUT)
            dpuVapiEsxcliNonTaskList.append(dpuVapiEsxcliNonTask)
         return dpuVapiEsxcliNonTaskList
      else:
         dpuEsxcliTasks = []

         # Keep a copy/backup of self._command as we need it for multiple DPUs
         commandBackup = self._command
         for dpu, depotUrl in itertools.zip_longest(self._dpus,
                                                    self._localDepots,
                                                    fillvalue=None):
            self._addExtraCommandArgs(depotUrl=depotUrl, dpu=dpu)
            dpuEsxcliTask = DpuEsxcliTask(dpu, self._command, self._isLocalcli,
                                          self._isInternal,
                                          timeout=self.TIMEOUT)
            self._command = commandBackup
            dpuEsxcliTasks.append(dpuEsxcliTask)
         return dpuEsxcliTasks

   def finalizeOperation(self, nativeOutput=True):
      """API to finalize the operation. Categorize the DPUs based on the
         compliance status.
      """
      for task in self._dpuTasks:
         self._results[task._dpuId] = {
            'status': task._status,
            'output': task.getOutput(native=nativeOutput)
         }

   def getResults(self):
      """Returns results of the operation.
      """
      return self._results

class VibListOperation(EsxcliOperation):
   """ Run vib list command on all DPUs.
   """
   USE_DEPOTS = False

   def __init__(self, isLocalcli=False):
      """ Construct a VibListOperation.
      """
      super(VibListOperation, self).__init__('vib list', isLocalcli)

   def getResults(self):
      """ Process self._results.
      """
      return self._results


class ComponentApplyEsxcliOperation(EsxcliOperation):
   """Run component apply command on all DPUs.
   """
   TIMEOUT = 800

   def __init__(self, compSpecs, depots, dryRun=False, noSigCheck=False):
      """Constructor.

         Parameters:
            compSpecs: a list of Component specification strings;
                       the format is <name> or <name>:<version>
            dryRun: whether to set --dry-run option
            noSigCheck: whether to set --no-sig-check option
      """
      cmd = 'component apply'
      extraArgs = dict()
      if dryRun:
         extraArgs['--dry-run'] = ""
      if noSigCheck:
         extraArgs['--no-sig-check'] = ""
      compSpecStr = ' '.join(['-n %s' % spec for spec in compSpecs])
      if USE_ESXCLI_VAPI:
         if compSpecs:
            # As map can't have duplicate entries making the first
            # name and 2nd -n onwards as a string value to key '-n'.
            extraArgs['-n'] = compSpecStr[3:]
      else:
         cmd = cmd + ' ' + compSpecStr

      super().__init__(cmd, depots=depots,
                       extraArgs=extraArgs)


class ComponentRemoveOperation(EsxcliOperation):
   """Run component remove command on all DPUs.
   """
   TIMEOUT = 800
   USE_DEPOTS = False

   def __init__(self, compSpecs, dryRun=False):
      """Constructor.

         Parameters:
            compSpecs: a list of Component specification strings;
                       the format is <name> or <name>:<version>
            dryRun: whether to set --dry-run option
      """
      cmd = 'component remove'
      extraArgs = dict()
      if dryRun:
         extraArgs['--dry-run'] = ""
      compSpecStr = ' '.join(['-n %s' % spec for spec in compSpecs])
      if USE_ESXCLI_VAPI:
         # As map can't have duplicate entries making the first
         # name and 2nd -n onwards as a string value to key '-n'.
         extraArgs['-n'] = compSpecStr[3:]
      else:
         cmd = cmd + ' ' + compSpecStr
      super().__init__(cmd, extraArgs=extraArgs)

class GetAndListCliOperation(EsxcliOperation):
   """ Handles get, list and signature verify esxcli command
       on all DPUs.
   """
   USE_DEPOTS = False
   BASEIMAGE_GET_CMD = 'baseimage get'
   ADDON_GET_CMD = 'addon get'
   PROFILE_GET_CMD = 'profile get'
   COMP_LIST_CMD = 'component list'
   COMP_SIG_VERIFY_CMD = 'component signature verify'
   VIB_LIST_CMD = 'vib list'
   VIB_GET_CMD = 'vib get'
   VIB_SIG_VERIFY_CMD = 'vib signature verify'

   def __init__(self, cmd, rebootingImage=False):
      """ Constructor.

         Parameters:
            cmd: Esxcli command to be executed.
            rebootingImage: If true, Displays information for image
                            which becomes active after a reboot, or
                            nothing if the pending-reboot image has
                            not been created yet.
      """
      extraArgs =  dict()
      if rebootingImage:
         extraArgs['--rebooting-image'] = ""
      super(GetAndListCliOperation, self).__init__(cmd, extraArgs=extraArgs)

class GetAndListWithSpecCliOperation(EsxcliOperation):
   """ Handles get and list commands with component or vib specs
       on all DPUs.
   """
   USE_DEPOTS = False
   COMP_VIB_LIST_CMD = 'component vib list'
   COMP_GET_CMD = 'component get'

   def __init__(self, cmd, specs=None, rebootingImage=False):
      """ Constructor.

         Parameters:
            cmd: Esxcli command to be executed.
            specs: A list of Component/Vib specification strings;
                   the format is <name> or <name>:<version>
            rebootingImage: If true, Displays information for image
                            which becomes active after a reboot, or
                            nothing if the pending-reboot image has
                            not been created yet.
      """
      extraArgs = dict()
      if rebootingImage:
         extraArgs['--rebooting-image'] = ""
      if specs:
         compSpecStr = _optionJoin(' -n ', specs)
         if USE_ESXCLI_VAPI:
            # As map can't have duplicate entries making the first
            # name and 2nd -n onwards as a string value to key '-n'.
            extraArgs['-n'] = compSpecStr[4:]
         else:
            cmd += compSpecStr
      super(GetAndListWithSpecCliOperation, self).__init__(cmd,
                                                           extraArgs=extraArgs)

class ProfileInstallUpdateOperation(EsxcliOperation):
   """Run profile install/update command on all DPUs.
   """
   TIMEOUT = 600

   def __init__(self, profileName, depots, isInstallCmd=True, dryRun=False,
                force=False, noLiveInstall=False, checkMaintMode=True,
                okToRemove=False, allowDowngrades=False, noSigCheck=False,
                noHardwareWarn=False):
      """Constructor.

         Parameters:
            profileName: image profile name
            depot: a list of URL/local path to the depots
            isInstallCmd: True for profile install command, otherwise profile
                          update command
            dryRun: whether to set --dry-run option
            force: whether to set --force option
            noLiveInstall: whether to set --no-live-install option
            checkMaintMode: reverse flag, --maintenance-mode option will be
                            added if unset
            okToRemove: whether to set --ok-to-remove option, only applicable
                        to profile install command
            allowDowngrades: whether to set --allow-downgrades option, only
                             applicable to profile update command
            noSigCheck: whether to set --no-sig-check option
            noHardwareWarn: whether to set --no-hardware-warning option
      """
      extraArgs = {
         '-p': profileName,
      }
      if dryRun:
         extraArgs['--dry-run'] = ""
      if force:
         extraArgs['--force'] = ""
      if noLiveInstall:
         extraArgs['--no-live-install'] = ""
      if not checkMaintMode:
         extraArgs['--maintenance-mode'] = ""
      if okToRemove:
         if not isInstallCmd:
            raise ValueError('okToRemove can only be set with profile install '
                             'command')
         extraArgs['--ok-to-remove'] = ""
      if allowDowngrades:
         if isInstallCmd:
            raise ValueError('allowDowngrades can only be set with profile '
                             'update command')
         extraArgs['--allow-downgrades'] = ""
      if noSigCheck:
         extraArgs['--no-sig-check'] = ""
      if noHardwareWarn:
         extraArgs['--no-hardware-warning'] = ""

      super().__init__('profile install' if isInstallCmd else 'profile update',
                       depots=depots, extraArgs=extraArgs)


class ProfileValidateOperation(EsxcliOperation):
   """Run profile validate command on all DPUs.
   """
   def __init__(self, profileName, depot):
      """Constructor.

         Parameters:
            profileName: image profile name
            depot: a list of URL/local path to the depots
      """
      extraArgs = {
         '-p': profileName,
      }

      super().__init__('profile validate', depots=depot,
                       extraArgs=extraArgs)


class ApplyEsxcliOperation(EsxcliOperation):
   """Operation of apply command on DPUs.
   """
   TIMEOUT = 600

   def __init__(self, specDict, depots, componentsToDownload, dryRun=False,
                noLiveInstall=False, noSigCheck=False, noHardwareWarn=False):
      extraArgs = dict()
      if dryRun:
         extraArgs['--dry-run'] = ""
      if noLiveInstall:
         extraArgs['--no-live-install'] = ""
      if noSigCheck:
         extraArgs['--no-sig-check'] = ""
      if noHardwareWarn:
         extraArgs['--no-hardware-warning'] = ""
      super().__init__('apply', depots=depots, extraArgs=extraArgs)
      self._specDict = specDict
      self._componentsToDownload = componentsToDownload

   def _addExtraCommandArgs(self, dpu=None, depotUrl=None):
      """Add extra command arguments before task creation.
      """
      super()._addExtraCommandArgs(depotUrl=depotUrl)
      # Host JSON image spec file and append -s option.
      specUrl = HostDepot.hostTextFile(
            json.dumps(self._specDict), 'apply-swspec.json', dpu)
      if USE_ESXCLI_VAPI:
         self._extraArgs['-s'] = specUrl
      else:
         self._command += ' -s ' + specUrl

class VibInstallUpdateOperation(EsxcliOperation):
   """Operation of vib install/update command on all DPUs.
   """
   TIMEOUT = 600

   def __init__(self, vibUrls, vibNames, depots, isInstallCmd=True,
                dryRun=False, force=False, noLiveInstall=False,
                checkMaintMode=True, noSigCheck=False):
      """Constructor.

         Parameters:
            vibUrls: a list of URL/local path to VIBs
            vibNames: VIB name or name:version specifications
            depots: a list of URL/local path to the depots
            isInstallCmd: True for vib install command, otherwise vib update
                          command
            dryRun: whether to set --dry-run option
            force: whether to set --force option
            noLiveInstall: whether to set --no-live-install option
            checkMaintMode: reverse flag, --maintenance-mode option will be
                            added if unset
            noSigCheck: whether to set --no-sig-check option
      """
      extraArgs = dict()
      if dryRun:
         extraArgs['--dry-run'] = ""
      if force:
         extraArgs['--force'] = ""
      if noLiveInstall:
         extraArgs['--no-live-install'] = ""
      if not checkMaintMode:
         extraArgs['--maintenance-mode'] = ""
      if noSigCheck:
         extraArgs['--no-sig-check'] = ""

      if not vibUrls and not depots:
         raise ValueError('vibUrls and depots cannot be both empty')

      cmd = 'vib install' if isInstallCmd else 'vib update'
      if vibNames:
         vibSpecStr = _optionJoin(' -n ', vibNames)
         if USE_ESXCLI_VAPI:
            # As map can't have duplicate entries making the first
            # name and 2nd -n onwards as a string value to key '-n'.
            extraArgs['-n'] = vibSpecStr[4:]
         else:
            cmd += vibSpecStr

      super().__init__(cmd, depots=depots, extraVibs=vibUrls,
                       extraArgs=extraArgs)

   def _addExtraCommandArgs(self, dpu=None, depotUrl=None):
      """Add extra command arguments before task creation.
      """
      super()._addExtraCommandArgs(depotUrl=depotUrl)
      if self._extraVibs:
         # Add "-n name:version" for standalone VIBs.
         extraVibSpecs = self._hostDepot.extraVibSpecs
         vibSpecStr = _optionJoin(' -n ', extraVibSpecs)
         if extraVibSpecs:
            if USE_ESXCLI_VAPI:
               # As map can't have duplicate entries making the first
               # name and 2nd -n onwards as a string value to key '-n'.
               if hasattr(self._extraArgs,'-n'):
                  self._extraArgs['-n'] += vibSpecStr
               else:
                  self._extraArgs['-n'] = vibSpecStr[4:]
            else:
               self._command += vibSpecStr


class VibRemoveOperation(EsxcliOperation):
   """Operation of vib remove command on all DPUs.
   """
   TIMEOUT = 600
   USE_DEPOTS = False

   def __init__(self, vibNames, dryRun=False, force=False, noLiveInstall=False,
                checkMaintMode=True):
      """Constructor.

         Parameters:
            vibNames: VIB name or name:version specifications
            dryRun: whether to set --dry-run option
            force: whether to set --force option
            noLiveInstall: whether to set --no-live-install option
            checkMaintMode: reverse flag, --maintenance-mode option will be
                            added if unset
      """
      cmd = 'vib remove'
      vibSpecStr = _optionJoin(' -n ', vibNames)
      extraArgs = dict()
      if dryRun:
         extraArgs['--dry-run'] = ""
      if force:
         extraArgs['--force'] = ""
      if noLiveInstall:
         extraArgs['--no-live-install'] = ""
      if not checkMaintMode:
         extraArgs['--maintenance-mode'] = ""
      if USE_ESXCLI_VAPI:
         # As map can't have duplicate entries making the first
         # name and 2nd -n onwards as a string value to key '-n'.
         extraArgs['-n'] = vibSpecStr[4:]
      else:
         cmd += vibSpecStr
      super().__init__(cmd, extraArgs=extraArgs)


class AcceptanceGetSetOperation(EsxcliOperation):
   """Operation of acceptance get/set command on DPUs.
   """
   USE_DEPOTS = False

   def __init__(self, newLevel=None):
      """Constructor.

         Parameters:
            newLevel: if set, run set command with this acceptance level, see
                      Vib.ArFileVib.ACCEPTANCE_LEVELS for valid values;
                      otherwise, run get command.
      """
      extraArgs = {'--level': newLevel} if newLevel else {}
      cmd = 'acceptance set' if newLevel else \
            'acceptance get'
      super().__init__(cmd, extraArgs=extraArgs)


def removeSolutionOnDPUs(solutionNames, parentTask):
   """ Remove solutions on DPUs.

      solutionNames: The solutions to be removed.
      parentTask: The parent task for ESXi host solution removal.

      Return: NormalExit on success; otherwise, DpuOperationError.
   """
   try:
      removeSolOp = RemoveSolutionOperation(solutionNames, parentTask)
      if removeSolOp.run():
         return NormalExit()
      else:
         return DpuOperationError('Failed to remove solutions.')
   except DpuInfoError as e:
      log.exception('Failed to get DPU info: %s', e)
      if parentTask:
         notif = createNotification(DPUInfoError, DPUInfoError, [], [], ERROR)
         parentTask.updateNotifications([notif])
      return DpuOperationError('Failed to remove solutions on DPUs: %s' %
                               str(e))
   except Exception as e:
      msg = 'Failed to remove solutions on DPUs: %s' % str(e)
      log.exception(msg)
      return DpuOperationError(msg)

def vapiListInstalledOnDpus(isVibList, orphanVibs=False):
   """ List components or vibs on DPUs.

       isVibList: Flag to create compoent or vib list operation.
       orphanVibs: Flag to get orphan vibs list from Dpu.

       Return: NormalExit and results if succeeded,
               or DpuOperationError with no result.
   """
   _getOpName = 'vibs' if isVibList else 'components'

   try:
      listOp = VapiListOperation(isVibList, orphanVibs)
      if listOp.run():
         return NormalExit(), listOp._dpuResults
      else:
         return DpuOperationError('Failed to list %s on Dpu.' %
                                              _getOpName), None
   except DpuInfoError as e:
      log.exception('Failed to get DPU info: %s', str(e))
      return DpuOperationError('Failed to list %s on DPUs: %s' %
                               (_getOpName, str(e))) , None
   except Exception as e:
      msg = 'Failed to list %s on DPUs: %s' % (_getOpName, str(e))
      log.exception(msg)
      return DpuOperationError(msg), None

def _runEsxcliOperation(cmdName, opCls, *args, **kwargs):
   """Executes an esxcli operation, returns NormalExit and results if succeeded,
      or DpuOperationError when failed with no result.
   """
   try:
      op = opCls(*args, **kwargs)
      if not op.run():
         return (DpuOperationError('Failed to start esxcli operation on '
                    'DPU(s), see logs for details'),
                 op.getResults())

      results = op.getResults()

      LogLargeBuffer('esxcli %s results on DPUs: %s' % (cmdName, results),
                     log.debug)
      errStr = ''
      for dpu, result in results.items():
         if result['status'] != 0:
            errStr += ('esxcli %s failed on DPU %s:\n%s\n' % (
               cmdName, dpu, result['output']))
      if errStr:
         return DpuOperationError(errStr.rstrip('\n')), None
      else:
         return NormalExit(), results
   except DpuInfoError as e:
      log.exception('Failed to get DPU info')
      return (DpuOperationError('esxcli %s failed on DPU(s): %s'
                                % (cmdName, str(e))), None)
   except Exception as e:
      log.exception('Unexpected error when running esxcli %s on DPUs', cmdName)
      return (DpuOperationError('esxcli %s failed on DPU(s): %s'
                                % (cmdName, str(e))), None)

def componentsApplyOnDPUs(components, depots, esxcli=False, dryRun=False,
                          noSigCheck=False):
   """Apply the provided components from the depot bundle on all DPUs.

      components: Component name or name/version list. For name/version, they
                  are cat'ed with semicollon to a single string.
      depots: A depot file/URL list.
      esxcli: if set, call with esxcli rather than VAPI.
      dryRun/noSigCheck: whether to use --dry-run/--no-sig-check options, esxcli
                         only.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   if esxcli:
      return _runEsxcliOperation(
         'component apply', ComponentApplyEsxcliOperation, components, depots,
         dryRun=dryRun, noSigCheck=noSigCheck)

   if dryRun or noSigCheck:
      raise ValueError('--dry-run and --no-sig-check can only be used with '
                       'esxcli')

   try:
      op = ComponentApplyOperation(components, depots)
      if op.run():
         return NormalExit(), None
      else:
         return DpuOperationError('Failed to apply components'), None
   except DpuInfoError as e:
      msg = 'Failed to get DPU info: %s' % str(e)
      log.exception(msg)
      return DpuOperationError(msg), None
   except Exception as e:
      msg = 'Cannot apply components on dpu: %s' % str(e)
      log.exception(msg)
      return DpuOperationError(msg), None

def scanOnDPUs(softwareSpec, depots, relatedComponents, parentTask=None):
   """Scan DPUs.

      softwareSpec: The desired image spec.
      depots: A depot file/URL list.
      relatedComponents: The (name, version) pairs of components related
                         to the desired image.
      parentTask: The parent task for ESXi host scan.

      Return: The DPU scan results or None.
   """
   try:
      scanOp = ScanOperation(softwareSpec, depots,
                             relatedComponents, parentTask)
      dpuScanResult = scanOp.run()
      if dpuScanResult is None:
         raise DpuOperationError('Failed to perform scan task on DPU')
      return dpuScanResult
   except DpuInfoError as e:
      log.error('Failed to get DPU info: %s', e)
      if parentTask:
         notif = createNotification(DPUInfoError, DPUInfoError, [], [], ERROR)
         parentTask.updateNotifications([notif])
   except Exception as e:
      log.error('Cannot scan dpu images: %s', str(e))
      raise
   return None

def applyOnDPUs(softwareSpec, depots, componentsToDownload, esxcli=False,
                parentTask=None, dryRun=False, noLiveInstall=False,
                noSigCheck=False, noHardwareWarn=False):
   """ Apply desired image on DPUs.

      softwareSpec: The desired image spec.
      depots: A depot file/URL list.
      componentsToDownload: Related componnets to be downloaded.
      esxcli: if set, call with esxcli rather than VAPI.
      parentTask: VAPI only, the parent task for ESXi host apply.
      dryRun/noLiveInstall/noSigCheck/noHardwareWarn: esxcli only, set
         --dry-run, --no-live-install, --no-sig-check and --no-hardware-warning
         option respectively.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result;
              for VAPI currently always return None for results.
   """
   if esxcli:
      # esxcli invocation.
      return _runEsxcliOperation(
         'apply', ApplyEsxcliOperation, softwareSpec, depots,
         componentsToDownload, dryRun=dryRun, noLiveInstall=noLiveInstall,
         noSigCheck=noSigCheck, noHardwareWarn=noHardwareWarn)

   if dryRun or noLiveInstall or noSigCheck or noHardwareWarn:
      raise ValueError('VAPI apply does not support dryRun, noLiveInstall, '
            'noSigCheck or noHardwareWarn')

   try:
      applyOp = ApplyOperation(softwareSpec, depots, componentsToDownload,
                               parentTask)
      if applyOp.run():
         return NormalExit(), None
      else:
         return DpuOperationError('Desired image apply failed.'), None
   except DpuInfoError as e:
      log.exception('Failed to get DPU info: %s', e)
      if parentTask:
         notif = createNotification(DPUInfoError, DPUInfoError, [], [], ERROR)
         parentTask.updateNotifications([notif])
      return DpuOperationError('Desired image apply failed: %s' % str(e)), None
   except Exception as e:
      log.exception('Cannot apply dpu images: %s', str(e))
      return DpuOperationError('Desired image apply failed: %s' % str(e)), None

def componentRemoveOnDPUs(components, dryRun=False):
   """Run esxcli component remove on DPUs.

      Parameters:
         components: a list of Component specification strings;
                     the format is <name> or <name>:<version>
         dryRun: whether to set --dry-run option

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(
      'component remove', ComponentRemoveOperation, components,
       dryRun=dryRun)

def profileGetOnDPUs(rebootingImage=False):
   """List image profiles on DPUs.

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.PROFILE_GET_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.PROFILE_GET_CMD,
                              rebootingImage)

def profileInstallOnDPUs(profileName, depots, dryRun=False, force=False,
                         noLiveInstall=False, checkMaintMode=True,
                         okToRemove=False, noSigCheck=False,
                         noHardwareWarn=False):
   """Run esxcli profile install on DPUs.

      Parameters:
         profileName: image profile name
         depot: a list of URL/local path to the depots
         dryRun: whether to set --dry-run option
         force: whether to set --force option
         noLiveInstall: whether to set --no-live-install option
         checkMaintMode: reverse flag, --maintenance-mode option will be
                         added if unset
         okToRemove: whether to set --ok-to-remove option
         noSigCheck: whether to set --no-sig-check option
         noHardwareWarn: whether to set --no-hardware-warning option

      Returns: NormalExit and results if succeeded,
               or DpuOperationError with no result.
   """
   return _runEsxcliOperation(
      'profile install', ProfileInstallUpdateOperation, profileName, depots,
      isInstallCmd=True, dryRun=dryRun, force=force,
      noLiveInstall=noLiveInstall, checkMaintMode=checkMaintMode,
      okToRemove=okToRemove, noSigCheck=noSigCheck,
      noHardwareWarn=noHardwareWarn)

def profileUpdateOnDPUs(profileName, depots, dryRun=False, force=False,
                        noLiveInstall=False, checkMaintMode=True,
                        allowDowngrades=False, noSigCheck=False,
                        noHardwareWarn=False):
   """Run esxcli profile update on DPUs.

      profileName: image profile name
      depot: a list of URL/local path to the depots
      dryRun: whether to set --dry-run option
      force: whether to set --force option
      noLiveInstall: whether to set --no-live-install option
      checkMaintMode: reverse flag, --maintenance-mode option will be
                      added if unset
      allowDowngrades: whether to set --allow-downgrades option
      noSigCheck: whether to set --no-sig-check option
      noHardwareWarn: whether to set --no-hardware-warning option

      Returns: NormalExit and results if succeeded,
               or DpuOperationError with no result.
   """
   return _runEsxcliOperation(
      'profile update', ProfileInstallUpdateOperation, profileName, depots,
      isInstallCmd=False, dryRun=dryRun, force=force,
      noLiveInstall=noLiveInstall, checkMaintMode=checkMaintMode,
      allowDowngrades=allowDowngrades, noSigCheck=noSigCheck,
      noHardwareWarn=noHardwareWarn)

def profileValidateOnDPUs(profileName, depot):
   """Run esxcli profile validate on DPUs.

      profileName: image profile name
      depot: a list of URL/local path to the depots

      Returns:  NormalExit and results if succeeded,
                or DpuOperationError with no result.
   """
   return _runEsxcliOperation('profile validate', ProfileValidateOperation,
                              profileName, depot)

def vibInstallUpdateOnDPUs(vibUrls, vibNames, depots, isInstallCmd=True,
                           dryRun=False, force=False, noLiveInstall=False,
                           checkMaintMode=True, noSigCheck=False):
   """Run esxcli vib install/update on DPUs.

      Parameters:
         vibUrls: a list of URL/local path to VIBs
         vibNames: VIB name or name:version specifications
         depots: a list of URL/local path to the depots
         isInstallCmd: True for vib install command, otherwise vib update
                       command
         dryRun: whether to set --dry-run option
         force: whether to set --force option
         noLiveInstall: whether to set --no-live-install option
         checkMaintMode: reverse flag, --maintenance-mode option will be
                         added if unset
         noSigCheck: whether to set --no-sig-check option
   """
   return _runEsxcliOperation(
      'vib install' if isInstallCmd else 'vib update',
      VibInstallUpdateOperation, vibUrls, vibNames, depots,
      isInstallCmd=isInstallCmd, dryRun=dryRun, force=force,
      noLiveInstall=noLiveInstall, checkMaintMode=checkMaintMode,
      noSigCheck=noSigCheck)

def vibRemoveOnDPUs(vibNames, dryRun=False, force=False, noLiveInstall=False,
                    checkMaintMode=True):
   """Run esxcli vib remove on DPUs.

      Parameters:
         vibNames: VIB name or name:version specifications
         dryRun: whether to set --dry-run option
         force: whether to set --force option
         noLiveInstall: whether to set --no-live-install option
         checkMaintMode: reverse flag, --maintenance-mode option will be
                         added if unset
   """
   return _runEsxcliOperation(
      'vib remove', VibRemoveOperation, vibNames, dryRun=dryRun, force=force,
      noLiveInstall=noLiveInstall, checkMaintMode=checkMaintMode)

def baseImageGetOnDPUs(rebootingImage=False):
   """Display the installed baseimage on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.BASEIMAGE_GET_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.BASEIMAGE_GET_CMD,
                              rebootingImage)

def addonGetOnDPUs(rebootingImage=False):
   """Display the installed Addon on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.ADDON_GET_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.ADDON_GET_CMD,
                              rebootingImage)

def acceptanceGetOnDPUs():
   """Run esxcli acceptance get on DPUs, returns an exit state and DPU results.
   """
   return _runEsxcliOperation('acceptance get', AcceptanceGetSetOperation)

def acceptanceSetOnDPUs(newLevel):
   """Run esxcli acceptance set on DPUs, returns an exit state and DPU results.
      Parameters:
         newLevel: new acceptance level, see Vib.ArFileVib.ACCEPTANCE_LEVELS for
                   valid values.
   """
   return _runEsxcliOperation('acceptance set', AcceptanceGetSetOperation,
                              newLevel=newLevel)

def compGetOnDPUs(compSpecs=None, rebootingImage=False):
   """List image profiles on DPUs.

      Parameters:
         compSpecs: a list of Component specification strings;
                     the format is <name> or <name>:<version>
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListWithSpecCliOperation.COMP_GET_CMD,
                              GetAndListWithSpecCliOperation,
                              GetAndListWithSpecCliOperation.COMP_GET_CMD,
                              compSpecs, rebootingImage)

def vibGetOnDPUs(rebootingImage=False):
   """Get detailed information on all or mentioned vibs on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.VIB_GET_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.VIB_GET_CMD,
                              rebootingImage)

def compVibListOnDPUs(compSpecs=None, rebootingImage=False):
   """List vibs in a given component on DPU(s).

      Parameters:
         compSpecs: a list of Component specification strings;
                     the format is <name> or <name>:<version>
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListWithSpecCliOperation.COMP_VIB_LIST_CMD,
                              GetAndListWithSpecCliOperation,
                              GetAndListWithSpecCliOperation.COMP_VIB_LIST_CMD,
                              compSpecs, rebootingImage)

def compListOnDPUs(rebootingImage=False):
   """Display the component list on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.COMP_LIST_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.COMP_LIST_CMD,
                              rebootingImage)

def compSigVerifyOnDPUs(rebootingImage=False):
   """Display signature verification for the components on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.COMP_SIG_VERIFY_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.COMP_SIG_VERIFY_CMD,
                              rebootingImage)

def vibListOnDPUs(rebootingImage=False):
   """Display the vib list on DPU(s).

         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.VIB_LIST_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.VIB_LIST_CMD,
                              rebootingImage)

def vibSigVerifyOnDPUs(rebootingImage=False):
   """Display signature verification for the vibs on DPU(s).

      Parameters:
         rebootingImage: If true, Displays information for image
                         which becomes active after a reboot, or
                         nothing if the pending-reboot image has
                         not been created yet.

      Return: NormalExit and results if succeeded,
              or DpuOperationError with no result.
   """
   return _runEsxcliOperation(GetAndListCliOperation.VIB_SIG_VERIFY_CMD,
                              GetAndListCliOperation,
                              GetAndListCliOperation.VIB_SIG_VERIFY_CMD,
                              rebootingImage)
