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

""" This module implements an ImageRunnable subclasses that runs DPU vAPI tasks.
"""

import esxutils
import logging

from abc import abstractmethod
from com.vmware import cis_client
from com.vmware.esx.settings_daemon_client import (Components,
   DataProcessingUnitCompliance, Software, Solutions,
   TaskInfo)
from com.vmware.esx.software_client import InstalledComponents, InstalledVibs
from com.vmware.esx.settings_daemon import data_processing_unit_client as dpuc

from datetime import datetime
from threading import Thread

from vmware.runcommand import runcommand
from vmware.esxsoftwarecli import List, Struct
from vmware.vapi.lib.connect import get_connector
from vmware.vapi.stdlib.client.factories import StubConfigurationFactory
from vmware.vapi.bindings.converter import TypeConverter

log = logging.getLogger(__name__)

from ..Errors import DpuInfoError
from .DpuLib import getDpuAlias, getDpuID, getDpuTransportAddr, isDpuUnavailable
from .TaskBase import (ImageRunnable, STATE_EARLY_SUCCEEDED, STATE_EARLY_FAILED,
                       STATE_FAILED, STATE_MAX_RETRY, STATE_NEXT_PHASE,
                       STATE_PENDING, STATE_RUNNING, STATE_SUCCEEDED,
                       STATE_TIMEDOUT, STATE_TO_EARLY_SUCCEEDED,
                       STATE_TO_EARLY_FAILED, STATE_TO_NEXT_PHASE,
                       Workflow, WorkflowPhase)

from ..Utils import XmlUtils
from .Constants import *

etree = XmlUtils.FindElementTree()

APPLY_WORKFLOW_NAME = 'apply_workflow'
APPLY_TASK_NAME = 'apply_task'
REMOVE_TASK_NAME = 'remove_task'
SCAN_TASK_NAME = 'scan_task'
LIST_TASK_NAME = 'list'

GET_ACCESS_TOKEN_CMD = '/bin/getAccessToken -d %s -s ESXIO_DESIRE_STATE'

# Error code for XML parse failure.
PARSE_ERROR = 256

def _getVapiStub(addr, dpuId):
   connector = get_connector('http', 'json',
                              url='http://%s/lifecycle-api' % addr)
   from vmware.vapi.security.oauth import create_oauth_security_context

   status, token = runcommand(GET_ACCESS_TOKEN_CMD % dpuId)
   if status == 0:
      secCtx = create_oauth_security_context(token.decode())
   else:
      raise Exception('failed to get DPU access token')

   connector.set_security_context(secCtx)
   return StubConfigurationFactory.new_std_configuration(connector)

class _EsxcliOutput(object):
   """ A class to parse esxcli command output in XML format.
   """

   LIST = 'list'
   NAME = 'name'
   ROOT = 'root'
   STRING = 'string'
   STRUCT = 'structure'
   TYPE = 'type'
   TYPE_NAME = 'typeName'

   def __init__(self, output):
      """
         Members:
            _output: The original output XML string.
      """
      self._output = output

   def getNativeOutput(self):
      """ Generate native output.
      """

      def removeNamespaceFromTagName(xml, namespace):
         """ Remove namespace from tag name.
         """
         xml.tag = xml.tag.replace(namespace, '')
         for child in list(xml):
            removeNamespaceFromTagName(child, namespace)

      def parseXML(xml):
         """ Recursively parse XML tree.
         """
         if xml.tag == self.STRING:
            # Simple string node at root
            return xml.text
         if xml.tag == self.LIST:
            # A list
            ret = []
            elementType = xml.get(self.TYPE)
            for child in list(xml):
               value = (parseXML(child) if elementType == self.STRUCT
                        else child.text)
               ret.append(value)
            return List(elementType, ret)
         # A structure
         ret = dict()
         for child in list(xml):
            grandChild = list(child)[0]
            childName = child.get(self.NAME)
            if grandChild.tag == self.LIST:
               value = parseXML(grandChild)
            else:
               value = list(child)[0].text
            ret[childName] = value
         ret['structtype'] = xml.get(self.TYPE_NAME)
         return Struct(ret)

      try:
         xml = XmlUtils.ParseXMLFromString(self._output.encode())
      except etree.XMLSyntaxError as e:
         log.error('Failed to load ESXCLI XML output: %s', str(e))
         return None

      namespace = xml.tag[:xml.tag.find('}') + 1]
      if namespace:
         removeNamespaceFromTagName(xml, namespace)

      try:
         root = [c for c in list(xml) if c.tag == self.ROOT][0]
         root = list(root)[0]
      except IndexError:
         log.error('Cannot find output object in ESXCLI XML output')
         return None

      return parseXML(root)

class DpuTask(ImageRunnable):
   """ ImageRunnable subclass that run a task on a single DPU.
   """

   def __init__(self, taskName, dpuInfo, parentTask, maxRetry, timeout):
      """ Construct a DpuTask.

          Parameters:

          taskName: The task name that maps to the vAPI method.
          dpuInfo: The DPU info data.
          parentTask: The ESXi host task if any.
          maxRetry: The maximum retry count.
          timeout: The timeout period.
      """
      self._dpuInfo = dpuInfo
      try:
         self._dpuId = getDpuID(dpuInfo)
         self._ipAddr = getDpuTransportAddr(dpuInfo)
         self._dpuAlias = getDpuAlias(dpuInfo)
      except DpuInfoError as e:
         log.error('Failed to get DPU ID or IP: %s', str(e))
         if parentTask:
            self._parentTask = parentTask
            self.updateParentTaskNotification(DpuInfoError, [], ERROR)
         self._state = STATE_FAILED
         return

      super(DpuTask, self).__init__(taskName, self._dpuId, parentTask, maxRetry,
                                    timeout)

   def postProcess(self):
      pass

   def __del__(self):
      self.postProcess()

   @abstractmethod
   def updateTask(self):
      """ Check the task state; update state if the task succeeded or failed.
      """
      self.updateTask()

   # State relation graph for DpuTask.
   stateTransitionGraph = {STATE_FAILED: [STATE_PENDING, STATE_MAX_RETRY],
                           STATE_MAX_RETRY: None,
                           STATE_PENDING: [STATE_RUNNING, STATE_FAILED],
                           STATE_RUNNING: [STATE_SUCCEEDED, STATE_FAILED,
                                           STATE_TIMEDOUT],
                           STATE_SUCCEEDED: None,
                           STATE_TIMEDOUT: None,}

   # State transition map for DpuEsxcliTask.
   stateTransitionFuncMap = {STATE_PENDING: ImageRunnable.start,
                             STATE_RUNNING: updateTask,
                             STATE_FAILED: ImageRunnable.processFailed}

   # The long run state
   longRunState = [STATE_RUNNING]

class DpuEsxcliTask(DpuTask):
   """ ImageRunnable subclass that runs an ESXcli software command.

       dpuInfo: The info of the DPU where vAPI runs.
       command: The subcommand of esxcli software or softwareinternal.
       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.
       maxRetry: The maximum retry count.
       timeout: The timeout period.
   """

   ESXCLI_SOFTWARE = '/bin/esxcli --formatter=xml software'
   LOCALCLI_SOFTWARE = '/bin/localcli --formatter=xml software'
   INTERNAL_SOFTWARE = ('/bin/localcli --formatter=xml '
      '--plugin=/usr/lib/vmware/esxcli/int/ext/esxcli-softwareinternal.xml '
      'softwareinternal')
   SSH_CMD = ('/bin/ssh -q -o UserKnownHostsFile=/dev/null -o '
              'StrictHostKeyChecking=no')
   # User vpxuser is allowed in strict lockdown mode.
   DPU_SERVICE = '/bin/vim-cmd -U vpxuser combinersvc/dpu_services/'
   SSH_STATUS = DPU_SERVICE + 'get_status TSM-SSH %s'
   ENABLE_SSH = DPU_SERVICE + 'start TSM-SSH %s'
   DISABLE_SSH = DPU_SERVICE + 'stop TSM-SSH %s'

   ENABLE_SSH_FW = ['network', 'firewall', 'ruleset', 'client',
                    'add', '-r', 'sshClient']
   DISABLE_SSH_FW = ['network', 'firewall', 'ruleset', 'client', 'remove',
                     '-r', 'sshClient']

   POLL_INTERVAL = 1

   def __init__(self, dpuInfo, command, isLocalcli, isInternal, maxRetry=None,
                timeout=None):
      """ Construct a runnable object that runs an esxcli software command.
      """
      super(DpuEsxcliTask, self).__init__(command, dpuInfo, None,
                                          maxRetry, timeout)

      dest = 'root' + '@' + self._ipAddr
      prefix = self.ESXCLI_SOFTWARE
      if isLocalcli:
         prefix = (self.INTERNAL_SOFTWARE if isInternal
                                          else self.LOCALCLI_SOFTWARE)
      self._command = ' '.join([self.SSH_CMD, dest, prefix, command])
      self._child = None
      self._status = 0
      self._output = None
      self._nativeOutput = None

   def runCommand(self):
      """ Helper method to run the ssh command for this task.
      """
      sshAlreadyEnabled = False
      enabledSSH = False
      error = None
      try:
         status, output = runcommand(self.SSH_STATUS % self._dpuAlias)
         if status == 0:
            sshAlreadyEnabled = output.decode().strip() == 'Running'

         if not sshAlreadyEnabled:
            status, output = runcommand(self.ENABLE_SSH % self._dpuAlias)
            if status == 0:
               enabledSSH = True
            else:
               error = output.decode()
      except Exception as e:
         error = str(e)
      finally:
         if not (sshAlreadyEnabled or enabledSSH):
            self._status = 1
            self._output = 'Failed to enable ssh: %s.' % str(error)
            return

      enableFW = False
      try:
         esxutils.runCli(self.ENABLE_SSH_FW)
         enableFW = True
      except Exception as e:
         self._status = 1
         self._output = 'Failed to enable ssh firewall ruleset: %s.' % str(e)
         return

      try:
         log.info('Running command "%s" on DPU %s', self._command, self._dpuId)
         self._status, self._output = runcommand(self._command)
         self._output = self._output.decode()
      except Exception as e:
         self._status = 1
         self._output = 'Command failed with %s' % str(e)
      finally:
         if enableFW:
            try:
               esxutils.runCli(self.DISABLE_SSH_FW)
            except Exception as e:
               log.warning('Failed to disable ssh firewall ruleset: %s', str(e))

         if enabledSSH:
            try:
               status, output = runcommand(self.DISABLE_SSH % self._dpuAlias)
               if status != 0:
                  log.warning('Failed to disable ssh: %s', output.decode())
            except Exception as e:
               log.warning('Failed to disable ssh: %s', str(e))

   def _start(self):
      """ The private function to start the runnable.
      """
      try:
         self._child = Thread(target=DpuEsxcliTask.runCommand,
                             args=[self])
         self._child.start()
         self._state = STATE_RUNNING
      except Exception as e:
         log.error('Failed to start esxcli task on DPU %s: %s'
                   % (self._dpuId, str(e)))
         self._state = STATE_FAILED

   def _preprocessFailedState(self):
      """ Preprocess task before FAILED state transition:
            for serious issues such as LiveInstallationError, stop retry.
      """
      from ..Installer.InstallerCommon import FAILURE_WARNING
      if FAILURE_WARNING in self._output:
         self._maxTry = 1

   def updateTask(self):
      """ Check the task state; update state if the task succeeded or failed.
      """
      if not self._child.is_alive():
         if self._status != 0:
            self._preprocessFailedState()
            log.error('Command failed with error: %s', self._output)
            self._state = STATE_FAILED
         else:
            self._state = STATE_SUCCEEDED

   def getOutput(self, native=False):
      """ Return the output:
             If native is set,
                deserialization result, on deserialization succeed,
                XML output, on deserialization failure.
             If native is not set,
                XML output, on command succeed,
                error string, on command failure.
      """
      if native and self._state == STATE_SUCCEEDED:
         if not self._nativeOutput:
            outputParser = _EsxcliOutput(self._output)
            self._nativeOutput = outputParser.getNativeOutput()
            if self._nativeOutput is None:
               self._status = PARSE_ERROR
            else:
               return self._nativeOutput
      return self._output

class DpuVapiTask(DpuTask):
   """ ImageRunnable subclass that runs a DPU vAPI task.

       dpuInfo: The info of the DPU where vAPI runs.
       vapiCls: The vAPI client class.
       taskName: The task name that maps to the vAPI method.
       taskSpec: The argument for the vAPI call.
       maxRetry: The maximum retry count.
       timeout: The timeout period.
   """

   def __init__(self, dpuInfo, vapiCls, taskName, taskSpec, parentTask=None,
                dpuResults=None, maxRetry=None, timeout=None):
      """ Construct a runnable object that runs a vAPI task on DPU.
      """
      super(DpuVapiTask, self).__init__(taskName, dpuInfo, parentTask,
                                        maxRetry, timeout)

      self._vapiCls = vapiCls
      self._taskName = taskName
      self._taskSpec = taskSpec
      self._task = None
      self._taskId = None
      self._taskInfo = None
      self._progress = 0
      self._lastUpdateTime = datetime.min
      self._dpuResults = dpuResults

   def _start(self):
      """ The private function to start the runnable.
      """
      try:
         vapiSvc = self._vapiCls(_getVapiStub(self._ipAddr, self._dpuId))
         self._task = getattr(vapiSvc, self._taskName)(self._taskSpec)
         self._taskId = self._task.get_task_id()
         self._state = STATE_RUNNING
      except Exception as e:
         self.updateParentTaskNotification(TaskStartError,
                                           [self._name, self._dpuId],
                                           ERROR)
         log.error('Failed to start VAPI task on DPU %s: %s', self._dpuId,
                   str(e))
         self._state = STATE_FAILED

   def updateTask(self):
      """ Check the task state; update state if the task succeeded or failed.
      """
      try:
         taskSvc = cis_client.Tasks(_getVapiStub(self._ipAddr, self._dpuId))
         info = taskSvc.get(self._taskId)
         self._taskInfo = info.convert_to(TaskInfo)
         log.debug('Task %s status: %s', self._taskId, str(self._taskInfo))
         self.updateDpuResults()

         if self._taskInfo.notifications:
            dpuNotifs = self._taskInfo.notifications
            newNotifs = []
            newUpdateTime = self._lastUpdateTime
            infos = dpuNotifs.info or []
            warnings = dpuNotifs.warnings or []
            errors = dpuNotifs.errors or []
            for notif in infos + warnings + errors:
               if notif.time > self._lastUpdateTime:
                  newNotifs.append(notif)
                  if notif.time >= newUpdateTime:
                     newUpdateTime = notif.time
            if newNotifs:
               self._parentTask.updateNotifications(newNotifs)
            self._lastUpdateTime = newUpdateTime
      except Exception as e:
         self.updateParentTaskNotification(TaskUpdateError,
                                           [self._name, self._dpuId],
                                           ERROR)
         log.error('Failed to get task status: %s', str(e))
         self._state = STATE_FAILED

      if self._state == STATE_RUNNING:
         if self._taskInfo.status == STATE_SUCCEEDED:
            self.processSucceeded()
            log.info('DPU task %s on DPU %s succeeded.',
                     self._name, self._dpuId)
         elif self._taskInfo.status == STATE_FAILED:
            self._state = STATE_FAILED
            self.updateParentTaskNotification(TaskFailed,
                                              [self._name, self._dpuId],
                                              INFO)
            log.error('DPU task %s on DPU %s failed.',
                     self._name, self._dpuId)
         if self._taskInfo.progress:
            self._progress = self._taskInfo.progress.completed

   def updateDpuResults(self):
      """The API to update the DPU result to the DPU results.
      """
      pass

   def postProcess(self):
      """ Notify when all retries failed for UNAVAILABLE DPU
      """
      if not self.isSuccess() and isDpuUnavailable(self._dpuInfo):
         self.updateParentTaskNotification(UnavailableDpuTaskFailed,
                                           [self._name, self._dpuId],
                                           INFO)

class DpuVapiNonTask(DpuTask):
   """ ImageRunnable subclass that invokes DPU VAPI call that is not a task.

       dpuInfo: The info of the DPU where vAPI runs.
       vapiCls: The vAPI client class.
       taskName: The task name that maps to the vAPI method.
       dpuResults: The vib/component list task results for DPUs.
       param: The parameters required for vapi task.
   """

   def __init__(self, dpuInfo, vapiCls, taskName, dpuResults=None,
                                                      param=None):
      """ Construct a runnable object that runs a vAPI task on DPU.
      """
      super(DpuVapiNonTask, self).__init__(taskName, dpuInfo, None,
                                          maxRetry=None, timeout=None)

      self._vapiCls = vapiCls
      self._taskName = taskName
      self._dpuResults = dpuResults
      self._param = param
      self._dpuOutput = None
      self._nativeOutput = None
      self._child = Thread(target=self._makeVapiCall, args=())

   def _makeVapiCall(self):
      vibfilterSpec = None
      vapiSvc = self._vapiCls(_getVapiStub(self._ipAddr, self._dpuId))
      try:
         if self._param:
            vibfilterSpec = self._param['filterSpec']
            self._dpuOutput = getattr(vapiSvc, self._taskName)(vibfilterSpec)
         else:
            self._dpuOutput = getattr(vapiSvc, self._taskName)()
      except Exception as e:
         self._state = STATE_FAILED
         log.error('VAPI %s %s task failed on DPU %s: %s', self._vapiCls,
                   self._taskName, self._dpuId, str(e))

   def _start(self):
      """ The private function to start the runnable.
      """
      try:
         self._child.start()
         self._state = STATE_RUNNING
      except Exception as e:
         self._state = STATE_FAILED
         log.error('Failed to start VAPI task on DPU %s: %s', self._dpuId,
                   str(e))

   def updateTask(self):
      """ Check the dpu output and vapi timeout state.
          Update state if the task succeeded or failed.
      """
      if not self._child.is_alive():
         if self._state != STATE_FAILED:
            self._state = STATE_SUCCEEDED
            self.updateDpuResults()
         else:
            raise Exception('VAPI %s %s task failed on DPU %s.' %
                    (self._vapiCls, self._taskName, self._dpuId))

   def updateDpuResults(self):
      """The API to update the DPU results per managed DPU(s).
      """
      if self._dpuResults and self._dpuResults.get(self._dpuId):
         if self._dpuOutput:
            self._dpuResults[self._dpuId] = self._dpuOutput
         else:
            self._dpuResults[self._dpuId] = []

class DpuVapiEsxcliNonTask(DpuTask):
   """ Subclass of DpuTask, It runs esxcli commands on DPU.
   """
   ESXCLI_CMD = 'run_esxcli_cmd'
   ESXCLI_VAPI_CLS = dpuc.Esxcli if hasattr(dpuc, 'Esxcli') else None

   def __init__(self, dpuInfo, command, isLocalcli, isInternal,
                depots, extraArgs, maxRetry=None, timeout=None):
      """ Construct a runnable object that runs an esxcli software command.
      """
      super().__init__(self.ESXCLI_CMD, dpuInfo, None,
                       maxRetry, timeout)
      self._isLocalcli = isLocalcli
      self._isInternal = isInternal
      self._cmdStr = command
      self._depots = depots if depots else []

      # Make a copy of the dict as each DPU may have a different value for args
      self._extraArgs = extraArgs.copy() if extraArgs else {}
      self._status = 0
      self._output = None
      self._nativeOutput = None
      self._child = Thread(target=self._makeVapiCall, args=())

   def _makeVapiCall(self):
      vapiSvc = self.ESXCLI_VAPI_CLS(_getVapiStub(self._ipAddr, self._dpuId))
      try:
         dpuTask = getattr(vapiSvc, self.ESXCLI_CMD)
         cmd = self.ESXCLI_VAPI_CLS.Command
         cliTypeCls = self.ESXCLI_VAPI_CLS.CliType
         pairCls = self.ESXCLI_VAPI_CLS.Pair

         prefix = cliTypeCls.ESX_CLI
         if self._isLocalcli:
            if self._isInternal:
               prefix = cliTypeCls.INTERNAL_CLI
            else:
               prefix = cliTypeCls.LOCAL_CLI

         cmdType = cmdStrToEnum.get(self._cmdStr, None)
         if not cmdType:
            raise ValueError("%s is not a valid command." % self._cmdStr)

         argList = []
         for flag, val in self._extraArgs.items():
            if flag == '-n':
               # Creating a list of Pair(key, value) for each argument provided
               # in name string as required by the API.
               argList += [pairCls('-n', name.strip())
                           for name in val.split('-n ')]
               continue
            argList.append(pairCls(flag, val))

         dpuOutCls = dpuTask(cmd(prefix, cmdType, self._depots, argList))
         self._output = dpuOutCls.esxcli_output
         log.info('VAPI %s %s task started on DPU %s.',
                  self.ESXCLI_VAPI_CLS, cmdType, self._dpuId)
      except Exception as e:
         self._state = STATE_FAILED
         msg = "VAPI {0} {1} task failed on the DPU {2}: {3}".format(\
               self.ESXCLI_VAPI_CLS, cmdType, self._dpuId, str(e))
         self._output = msg
         log.error(msg)

   def _start(self):
      """ The private function to start the runnable.
      """
      try:
         self._child.start()
         self._state = STATE_RUNNING
      except Exception as e:
         self._state = STATE_FAILED
         log.error('Failed to start VAPI esxcli non-task on DPU %s: %s',
                   self._dpuId, str(e))

   def _preprocessFailedState(self):
      """ Preprocess task before FAILED state transition:
            for serious issues such as LiveInstallationError, stop retry.
      """
      from ..Installer.InstallerCommon import FAILURE_WARNING
      if FAILURE_WARNING in self._output:
         self._maxTry = 1

   def updateTask(self):
      """ Check the task state; update state if the task succeeded or failed.
      """
      if not self._child.is_alive():
         if self._status != 0:
            self._preprocessFailedState()
            log.error('Command failed with error: %s', self._output)
            self._state = STATE_FAILED
         else:
            self._state = STATE_SUCCEEDED

   def getOutput(self, native=False):
      """ Return the output:
         If native is set,
            deserialization result, on deserialization success,
            XML output, on deserialization failure.
         If native is not set,
            XML output, on command succeed,
            error string, on command failure.
      """
      if native and self._state == STATE_SUCCEEDED:
         if not self._nativeOutput:
            outputParser = _EsxcliOutput(self._output)
            self._nativeOutput = outputParser.getNativeOutput()
            if self._nativeOutput is None:
               self._status = PARSE_ERROR
            else:
               return self._nativeOutput
      return self._output

class ComponentApplyTask(DpuVapiTask):
   """ Subclass of DpuVapiTask that runs a component apply vAPI task.
   """
   TIMEOUT = 600

   def __init__(self, dpuInfo, components, depots):
      """ Constructor.

          dpuInfo: The info of the dup to run the vAPI.
          components: component name/version dict.
          depots: The depots that supports the apply.
      """
      applySpec = Components.ApplySpec(components=components,
                                       depots=depots)
      super(ComponentApplyTask, self).__init__(dpuInfo, Components,
                                               APPLY_TASK_NAME, applySpec)

class ScanTask(DpuVapiTask):
   """ Subclass of DpuVapiTask that runs an image scan vAPI task.
   """

   TIMEOUT = 100

   def __init__(self, dpuInfo, softwareSpec, depots, parentTask, dpuResults):
      """ Constructor.

          dpuInfo: The info of the dup to run the vAPI.
          softwareSpec: iThe desired image.
          depots: The depots that supports the apply.
          parentTask: The parent task for ESXi host scan.
          dpuResults: The scan task results for DPUs.
      """
      scanSpec = Software.ScanSpec(software_spec=softwareSpec,
                                   depots=depots)
      super(ScanTask, self).__init__(dpuInfo, Software, SCAN_TASK_NAME,
                                     scanSpec, parentTask, dpuResults)
      self._nativeResult = None

   def updateDpuResults(self):
      """The API method to update the DPU result into DPU results.

         Note: Scan returns HostCompliance from DPU now; Thus, a conversion is
               needed; when it returns DataProcessingUnitCompliance in the
               future, the conversion is redundant.
      """
      if self._taskInfo and self._taskInfo.result:
         self._nativeResult = TypeConverter.convert_to_python(
            self._taskInfo.result,
            DataProcessingUnitCompliance.get_binding_type())
         if self._dpuResults:
            self._dpuResults.compliance[self._dpuId] = self._nativeResult

class ScanWorkflowPhase(WorkflowPhase, ScanTask):
   """ The scan phase in the ApplyWithScan workflow.
   """

   def __init__(self, dpuInfo, softwareSpec, depots, parentTask):
      """ Constructor.

          dpuInfo: Info of DPU to be scanned.
          softwareSpec: The desired image.
          depots: Depot list used for scan.
          parentTask: The parent task if started from ESXi host scan.
          dpuResults: The results for all DPUs.
      """
      ScanTask.__init__(self, dpuInfo, softwareSpec, depots,
                        parentTask, dpuResults=None)

   def updateTask(self):
      """ Update the state of the scan phase as a task.
      """
      ScanTask.updateTask(self)
      if self._nativeResult:
         # TODO: add notification and process incompatible/unknown.
         complianceStatus = self._nativeResult.status
         if complianceStatus == COMPLIANT:
            self._state = STATE_TO_EARLY_SUCCEEDED
         elif complianceStatus == NON_COMPLIANT:
            self._state = STATE_TO_NEXT_PHASE
         else:
            self._state = STATE_TO_EARLY_FAILED

   # State transition graphh for scan phase.
   stateTransitionGraph = \
      WorkflowPhase.patchStateTransitionGraph(ScanTask.stateTransitionGraph)

   # State transition map for scan phase.
   stateTransitionFuncMap = {STATE_PENDING: ImageRunnable.start,
                             STATE_RUNNING: updateTask,
                             STATE_FAILED: ImageRunnable.processFailed}

class ApplySolutionTask(DpuVapiTask):
   """ Subclass of DpuVapiTask that runs solution apply vAPI task.
   """

   TIMEOUT = 800

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

          dpuInfo: The info of the dup to run the vAPI.
          solutions: The desired solutions.
          depots: The depots that supports the apply.
          parentTask: The parent task for ESXi host solution apply.
      """
      applySpec = Solutions.ApplySpec(solutions=solutions,
                                      depots=set(depots))
      super(ApplySolutionTask, self).__init__(dpuInfo, Solutions,
                                              APPLY_TASK_NAME, applySpec,
                                              parentTask)

class RemoveSolutionTask(DpuVapiTask):
   """ Subclass of DpuVapiTask that runs solution remove vAPI task.
   """
   TIMEOUT = 800

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

          dpuInfo: The info of the dup to run the vAPI.
          solutionNames: The solutions to be removed.
          parentTask: The parent task for ESXi host solution removal.
      """
      removeSpec = Solutions.RemoveSpec(solutions=set(solutionNames))
      super(RemoveSolutionTask, self).__init__(dpuInfo, Solutions,
                                               REMOVE_TASK_NAME, removeSpec,
                                               parentTask)

class ApplyTask(DpuVapiTask):
   """ Subclass of DpuVapiTask that runs an image apply vAPI task.
   """

   TIMEOUT = 800

   def __init__(self, dpuInfo, softwareSpec, depots, parentTask):
      """ Constructor.

          dpuInfo: The info of the dup to run the vAPI.
          softwareSpec: iThe desired image.
          depots: The depots that supports the apply.
          parentTask: The parent task if started from ESXi host scan.
      """
      applySpec = Software.ApplySpec(software_spec=softwareSpec,
                                     depots=depots)
      super(ApplyTask, self).__init__(dpuInfo, Software, APPLY_TASK_NAME,
                                      applySpec, parentTask)

class ApplyWorkflow(Workflow):
   """ ApplyWithScan workflow.
   """

   TIMEOUT = 1000

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

          dpuInfo: The info of the dup to run the vAPI.
          softwareSpec: The desired image.
          depots: The depots that supports the apply.
          parentTask: The parent task if started from ESXi host scan.
      """

      scanPhase = ScanWorkflowPhase(dpuInfo, softwareSpec, depots, parentTask)
      applyPhase = ApplyTask(dpuInfo, softwareSpec, depots, parentTask)
      workflowPhases = [scanPhase, applyPhase]
      try:
         self._dpuId = getDpuID(dpuInfo)
      except DpuInfoError as e:
         self.updateParentTaskNotification(DpuInfoError, [], ERROR)
         log.error('Failed to get DPU ID: %s', str(e))
         self._state = STATE_FAILED
         return

      super(ApplyWorkflow, self).__init__(workflowPhases, APPLY_WORKFLOW_NAME,
                                          self._dpuId, parentTask)

   # State relation graph for ApplyWithScan workflow.
   stateTransitionGraph = dict(DpuTask.stateTransitionGraph)
   stateTransitionGraph[STATE_EARLY_SUCCEEDED] = None
   stateTransitionGraph[STATE_EARLY_FAILED] = None

   # State transition map for ApplyWithScan workflow.
   stateTransitionFuncMap = {STATE_PENDING: ImageRunnable.start,
                             STATE_RUNNING: Workflow.updateWorkflow,
                             STATE_FAILED: ImageRunnable.processFailed}

   # The long run state.
   longRunState = [STATE_RUNNING]

class VapiDpuListTask(DpuVapiNonTask):
   """ Subclass of DpuVapiNonTask that runs a components
       or vibs list vAPI operations.
   """
   def __init__(self, dpuInfo, dpuResults, vibList, orphanVibs=False):
      """ Constructor.

          dpuInfo: The info of the dup to run the vAPI.
      """
      param = {}
      if orphanVibs:
         vibFilterSpec = InstalledVibs.FilterSpec(standalone_vibs_only=True)
         param['filterSpec'] = vibFilterSpec

      super(VapiDpuListTask, self).__init__(dpuInfo,
                                          InstalledVibs if vibList \
                                          else InstalledComponents,
                                          LIST_TASK_NAME,
                                          dpuResults,
                                          param)
