#!/usr/bin/python
#
# Copyright (c) 2014 Nutanix Inc. All rights reserved.
#
# Author: ssrinivasan@nutanix.com
#
# This script runs in ESX host after upgrade is initiated from CVM and CVM has
# shutdown. It basically runs the esx upgrade command on host and after reboot
# is done, removes host from maintenance mode. In case of any errors, it
# restarts the CVM service.

import functools
import glob
import logging
import os
import random
import re
import shutil
import string
import subprocess
import ssl
import sys
import time
import json

from tempfile import NamedTemporaryFile

from pyVim.connect import SmartConnect
from pyVmomi import vim, SoapAdapter, VmomiSupport

try:
  from httplib import HTTPException
except ImportError:
  from http.client import HTTPException

FILE = "/scratch/image.zip"
VIBS_DIR = "/scratch/nutanix_vibs"
DEPRECATED_VIBS_LIST_PATH = VIBS_DIR + '/deprecated_vib.json'
UPGRADE_MARKER = "/scratch/.host_upgrade"
CLEANUP_PENDING = "/scratch/.host_upgrade_cleanup"
MAINTENANCE_CLEANUP_PENDING = "/scratch/.host_maintenance_cleanup"
UPGRADE_PENDING = "/scratch/.host_upgrade_pending"
VIB_UPDATE_PENDING = "/scratch/.vib_update_pending"
UPGRADE_FAIL_MARKER = "/scratch/.host_upgrade_failed"
LOGFILE = "/scratch/log/esx_upgrade_script.log"
PCI_RECONFIG_MARKER = "/scratch/.cvm_pci_reconfig_marker"
NIC_CHECK_MARKER = "/scratch/.nic_check_marker"
REBOOT_MARKER = "/scratch/.reboot_marker"
MAINTENANCE_MODE_TIMEOUT = 30 * 60 # 30 minutes.
SOURCE_ESX_UPGRADE_FILE = "/scratch/esx_upgrade"
VMD_MARKER = "/scratch/.vmd_used_by_cvm"
ENABLE_DRIVERS = "/scratch/.enable_drivers"
SKIP_HARDWARE_WARNING_PATH = "/scratch/.skip_hardware_warning"
SKIP_HARDWARE_WARNING_FOR_SKYLAKE = "/scratch/.skip_hardware_warning_skylake"
NFS_VAAI_VIB = "nfs-vaai-plugin"
INTEL_OPTANE_VIB = "IntelOptanePMemMgmt"
DEL_PTAGENT_VIB = "dellptagent"
DEL_MNV_CLI_VIB= "BOSSN1-MNV-CLI"
UPGRADE_FINISHED = "/scratch/.upgrade_finished"
VIB_INSTALLED = "/scratch/.vib_installed"
CO_DATASTORE_UUID = "/bootbank/Nutanix/.ds_vmfs_uuid"
SOURCE_FACTORY_CONFIG_FILE = "/bootbank/Nutanix/factory_config.json"
SOURCE_HARDWARE_CONFIG_FILE = "/bootbank/Nutanix/hardware_config.json"
REMOVE_VIB_CMD = "esxcli software vib remove -n %s"
OPERATION_ID = "operationID"

def generate_op_id(id_len, prefix="ntnxlcm"):
  """
  Generates unique alphanumeric operation id

  Args:
    id_len (int): operation id (Op Id) length post prefix
    prefix (str, optional): prefix for Op Id, default is 'ntnxlcm'

  Returns:
    alphanumeric OpId (string)
    Eg: 'ntnxlcm-15svz3'
  """
  random_id_str = ''.join(random.choice(string.ascii_letters +
                                        string.digits) for _ in range(id_len))
  return '-'.join([prefix, random_id_str])

class OpId:
  """
  Context Manager that sets and restores unique operation id for
  each vmodl api call
  Operations performed:
  1. Gets unique operation id (op id)
  2. Before the api call, sets the op id in vmomi request context
  3. Post the api call, removes the current op id and
     restores the previous id in request context
  """

  def __init__(self, id_len=6):
    """
    Gets unique op id from generate_op_id of id_le size

    Args:
      id_len (int, optional): op id suffix length. Defaults to 6.
    """
    self.opId = generate_op_id(id_len)
    logging.info('Using VMODL operationId: %s' % self.opId)

  def __enter__(self):
    """
    Called before calling the api
    Sets op id in the request context
    """
    reqCtx = VmomiSupport.GetRequestContext()
    self.prevId = reqCtx.get(OPERATION_ID)
    reqCtx[OPERATION_ID] = self.opId
    return self

  def __exit__(self, *args):
    """
    Called post executing the api
    Removes op id for the api from request context
    """
    reqCtx = VmomiSupport.GetRequestContext()
    if self.prevId is not None:
      reqCtx[OPERATION_ID] = self.prevId
    else:
      del reqCtx[OPERATION_ID]

def mark_upgrade_failed(msg):
  """
  Mark upgrade operation as failed. Cleanup markers.
  """
  logging.error("Upgrade operation failed")

  # Create marker to signal failure.
  with open(UPGRADE_FAIL_MARKER, 'w') as fh:
    fh.write(msg)

  # Cleanup markers.
  for marker in [
      # /etc/rc.local.d/local.sh calls into esx_upgrade on reboot. If
      # UPGRADE_PENDING marker is not removed, it will try to re-trigger
      # the upgrade workflow without a profile name.
      UPGRADE_PENDING,
      # Remove reboot marker to avoid possible reboot loops.
      REBOOT_MARKER]:
    try:
      os.remove(marker)
      logging.info("Removed marker '%s'" % marker)
    except FileNotFoundError:
      # Ignore if marker is not present.
      pass

def cleanup_upgrade_pending():
  """
  Mark upgrade operation as not pending.
  """
  # Cleanup marker.
  os.remove(UPGRADE_PENDING) if os.path.exists(UPGRADE_PENDING) else None

def get_svm_dir(cvm_ds):
  """
  Returns the directory where SVM files are stored in the local datastore.
  """

  directory = '/vmfs/volumes/' + cvm_ds + '/ServiceVM_Centos/'

  logging.info("Returning SVM directory path %s" % directory)
  return directory

def get_svm_dir_co(uuid):
  """
  Returns the path to datastore on CO node
  """
  directory = '/vmfs/volumes/' + uuid
  logging.info("Returning CO datastore directory path %s" % directory)
  return directory

def in_maintenance_mode(host):
  """
  Returns true if host is in maintenance mode.
  """
  return host.summary.runtime.inMaintenanceMode

def exit_maintenance_mode(host):
  """
  Removes the Host out of maintenance mode.
  """
  if host.summary.config.product.version < "5.1.0":
    cmd = "vim-cmd hostsvc/maintenance_mode_exit"
    ret, out, err = run_command(cmd)
    logging.info("Maintenance mode exit triggered: ret %s out %s err %s" % (
                 ret, out, err))
    if ret:
      return False
    return True

  with OpId():
    task = host.ExitMaintenanceMode_Task(timeout=0)

  if task.info.state == vim.TaskInfo.State.error:
    logging.error("Task %s failed with error %s" % (task.info.name,
                                                    task.info.error.msg))
    return False
  else:
    logging.info("Exit maintenance mode ret: %s" % task.info.state)
    return True

def set_advanced_config_option(host, option_name, new_value):
  """
  Sets the value of the specified advanced option.
  """
  with OpId():
    option_array = host.configManager.advancedOption.QueryView(option_name)
  if not option_array:
    logging.error("Option '%s' was not found." % option_name)
    return False
  logging.info("Changing the value of '%s' from '%s' to '%s'" %
               (option_name, option_array[0].value, new_value))
  option_array[0].value = new_value
  with OpId():
    host.configManager.advancedOption.UpdateValues(changedValue=option_array)
  return True

def remove_nutanix_script_bootcfg():
  """
  Remove the nutanix.tgz entry added earlier from boot.cfg in /bootbank and
  /altbootbank directory of ESX.
  """
  esx_boot_dirs = [ "/bootbank", "/altbootbank" ]

  for boot_dir in esx_boot_dirs:
    boot_file = os.path.join(boot_dir, "boot.cfg")

    if not os.path.exists(boot_file):
      logging.info("Nutanix bootcfg file at %s does not exist" % boot_file)
      continue

    boot_cfg = None
    with open(boot_file, 'r') as fd:
      boot_cfg = fd.readlines()

    with NamedTemporaryFile(mode='w') as fd:
      for idx, line in enumerate(boot_cfg):
        if "modules=" in line:
          break

      line = boot_cfg[idx]
      if "nutanix.tgz" in line:
        index = line.find(" --- nutanix.tgz")
        line = line[:index] + line[index+len(" --- nutanix.tgz"):]
      else:
        continue
      boot_cfg[idx] = line

      fd.seek(0, os.SEEK_SET)
      fd.truncate()
      fd.write("".join(boot_cfg))
      fd.flush()
      try:
        shutil.copy(fd.name, boot_file)
        logging.info("Copied boot cfg to %s" % boot_dir)
      except Exception as e:
        logging.info("Exception raised while copying boot cfg file %s" % e)
        return False
  return True

def pynfs_supported(host):
  """
  Check if pynfs server is running on this system. We check if local NTNX
  datastore is NAS datastore or VMFS datastore. If NAS datastore, it should
  be mounted on pynfs server.
  """
  for ds in host.configManager.datastoreSystem.datastore:
    if "NTNX-local-ds" in ds.name:
      return isinstance(ds.info, vim.host.NasDatastoreInfo)
  return False

def remove_saved_esx_upgrade_script(target_file):
  """
  Removes the saved esx upgrade once upgrade operation completes
  """
  if target_file and os.path.exists(target_file):
    logging.info("Removing %s file" % target_file)
    os.remove(target_file)

def try_remove_saved_vmd_marker(target_dir):
  """
  Owner: Platforms Team
  Removes the saved vmd marker once upgrade operation completes, if exists.
  """
  vmd_backup = vmd_marker_backup_file(target_dir)
  if vmd_backup and os.path.exists(vmd_backup):
    logging.info("Removing %s file" % vmd_backup)
    os.remove(vmd_backup)

def post_upgrade_ops(host, target_dir, target_file):
  """
  Cleans up stuff which are needed post install. These include cleaning up
  Image bundle, moving host out of maintenance mode.
  """
  # Remove the saved esx_upgrade script
  remove_saved_esx_upgrade_script(target_file)
  try_remove_saved_vmd_marker(target_dir)

  logging.info("Removing upgrade file %s" % FILE)
  if os.path.exists(FILE):
    os.remove(FILE)

  logging.info("Removing Nutanix Vibs directory")
  if os.path.exists(VIBS_DIR):
    shutil.rmtree(VIBS_DIR)

  logging.info("Exiting maintenance mode")
  while in_maintenance_mode(host):
    ret = exit_maintenance_mode(host)
    if ret:
      # Exit maintenance mode task did not throw error. Wait for host to come
      # out of maintenance mode.
      while in_maintenance_mode(host):
        logging.info("Exiting maintenance mode in progress ... Waiting for "
                     "5 secs")
        time.sleep(5)
      break
    else:
      logging.error("Error while coming out of maintenance mode ... Retrying"
                    " in 5 secs")
      time.sleep(5)
  logging.info("Host is not in maintenance mode")

  def make_long(val):
    """Convert the given int to a long in both python 2.X and 3.X

    pyvmomi Requires a "long" type to update the config values. 3.X does not
    have type long, thus we must use the VmomiSupport.long. In 2.X, VmomiSupport
    does not have long.
    """
    try:
      return VmomiSupport.long(val)
    except AttributeError:
      return long(val)

  host_custom_config = [("NFS.HeartbeatTimeout", make_long(30)),
                        ("NFS.MaxVolumes", make_long(64)),
                        ("Net.TcpipHeapSize", make_long(32)),
                        ("SunRPC.MaxConnPerIP", make_long(64)),
                        ("UserVars.SuppressShellWarning", make_long(1))]

  for ii in host_custom_config:
    set_advanced_config_option(host, ii[0], ii[1])

  if host.summary.config.product.version >= "5.5":
    set_advanced_config_option(host, "Net.TcpipHeapMax", make_long(512))
  else:
    set_advanced_config_option(host, "Net.TcpipHeapMax", make_long(128))

  esx_boot_dirs = [ "/bootbank", "/altbootbank" ]
  # Remove the nutanix.tgz file from /bootbank and /altbootbank.
  for boot_dir in esx_boot_dirs:
    filename = os.path.join(boot_dir, "nutanix.tgz")
    if os.path.exists(filename):
      os.unlink(filename)

  if not remove_nutanix_script_bootcfg():
    logging.error("Unable to remove modification from boot script")

  if os.path.exists("/etc/rc.local.d/nutanix.sh"):
    if host.summary.config.product.version < "5.1.0":
      rc_file = "/etc/rc.local"
    else:
      rc_file = "/etc/rc.local.d/local.sh"

    rc_data = open(rc_file, 'r').read()
    nutanix_data = open("/etc/rc.local.d/nutanix.sh", 'r').read()
    # If contents are different copy nutanix.sh file to rc file.
    if rc_data != nutanix_data:
      logging.info("Update %s with nutanix contents" % rc_file)
      shutil.copy("/etc/rc.local.d/nutanix.sh", rc_file)
    else:
      logging.info("%s already has contents" % rc_file)
      # If local.sh file is same nutanix.sh remove nutanix.sh file
      if os.path.exists("/etc/rc.local.d/nutanix.sh"):
        os.unlink("/etc/rc.local.d/nutanix.sh")

  # Make sure /bootbank/pynfs.tgz and /bootbank/svmboot.tgz exist on
  # platform running pynfs local datastore.
  if not pynfs_supported(host):
    return

  for fl in [ "pynfs.tar.gz", "svmboot.tar.gz" ]:
    dest = os.path.join("/bootbank", fl)
    if not os.path.exists(dest):
      src = "/altbootbank/%s" % fl
      if os.path.exists(src):
        logging.info("Moving file %s to %s" % (src, dest))
        shutil.move(src, dest)
        continue
      src = "/bootbank/nutanix/%s" % fl
      if os.path.exists(src):
        logging.info("Moving file %s to %s" % (src, dest))
        shutil.move(src, dest)
        continue
      src = "/altbootbank/nutanix/%s" % fl
      if os.path.exists(src):
        logging.info("Moving file %s to %s" % (src, dest))
        shutil.move(src, dest)
        continue
  return

def maintenance_cleanup(host):
  """
  Perform cleanup when node has moved out of maintenance mode. Host is removed
  out of maintenance mode.
  """
  logging.info("Exiting maintenance mode")
  while in_maintenance_mode(host):
    ret = exit_maintenance_mode(host)
    if ret:
      # Exit maintenance mode task did not throw error. Wait for host to come
      # out of maintenance mode.
      while in_maintenance_mode(host):
        logging.info("Exiting maintenance mode in progress ... Waiting for "
                     "5 secs")
        time.sleep(5)
      break
    else:
      logging.error("Error while coming out of maintenance mode ... Retrying"
                    " in 5 secs")
      time.sleep(5)
  logging.info("Host is not in maintenance mode")
  return

def run_command(cmd):
  """
  Runs a command and returns output and error values.
  """
  pp = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                        shell=True)
  out, err = pp.communicate()

  if isinstance(out, bytes):
    out = out.decode()
  if isinstance(err, bytes):
    err = err.decode()
  return pp.returncode, out, err

def nutanix_cvm(vmlist):
  """
  Return the nutanix CVM given the vm list.
  """
  for vm in vmlist:
    if "ServiceVM_Centos.vmx" in vm.config.files.vmPathName:
      return vm
  return None

def check_nic_connectivity():
  """
  Check whether at least one of the NICs is up. If yes, return True.
  If no nics are up, and both ixgbe and ixgben modules are disabled, enable
  ixgben module, as upgrade form 6.0 to 6.5 disables ixgben module. Reboot the
  node after enabling the module.
  Returns:
      (check_passed, reboot_needed)
      check_passed (bool): True, proceed, False, error.
      reboot_needed (bool): True, reboot, False, no reboot necessary.
  """
  cmd = "esxcli network nic list"
  ret, out, err = run_command(cmd)
  logging.info("Executing cmd: %s, ret %s out %s err %s" % (cmd, ret, out,
                                                            err))
  if ret:
    logging.error("Unable to run cmd %s, ret %s out %s err %s" % (cmd, ret,
                  out, err))
    return False, False
  else:
    nics_link_pattern = re.compile("\S+\s+\S+\s+\S+\s+\S+\s+(\S+).*")
    nics_link_status_list = nics_link_pattern.findall(out)
    if "Up" in nics_link_status_list:
      logging.info("NICs have link up and associated with drivers.. No need "
                   "to enable ixgben driver")
      return True, False

  driver_list = ["ixgben", "ixgbe"]
  mod_enable_pattern = re.compile("\S+\s+\S+\s+(\S+).*")
  for driver in driver_list:
    disabled = False
    cmd = "esxcli system module list | grep %s" % driver
    ret, out, err = run_command(cmd)
    logging.info("Ran cmd %s, ret %s out %s err %s" % (cmd, ret, out, err))
    if ret == 0:
      match = mod_enable_pattern.search(out)
      if match:
        disabled = (match.group(1) == "false")
    # ixgbe/ixgben drivers are not disabled.
    if not disabled:
      logging.info("%r driver is enabled or not present, out %s"
                   % (driver, out))
      break
  else:
    # Both ixgbe and ixgben drivers are disabled.
    logging.info("Enable ixgben driver and reboot the host")
    cmd = "esxcli system module set -e true -m ixgben"
    ret, out, err = run_command(cmd)
    logging.info("Ran cmd %s, ret %s out %s err %s" % (cmd, ret, out, err))
    if ret:
      logging.error("Unable to run cmd %s, ret %s" % (cmd, ret))
      return False, False
    logging.info("Rebooting of host is required due to "
                 "change in network driver settings.")
    return True, True
  return True, False

def get_vib_name(vib_file):
  """
  Run excli command wih vib file to get vib info. Vib info contain vib name and
  other info.
  Returns:
    string : VIB name on success, None on failure.
  """
  # Execute command on host console to get vib_file info.
  vib_name = ''
  cmd = 'esxcli software sources vib get -v %s' % vib_file
  logging.info("Executing cmd: %s" % cmd)
  rv, out, err = run_command(cmd)
  if rv != 0:
    logging.error("Failed to run cmd %r on host, ret %d, stdout %s, stderr %s"
                  % (cmd, rv, out, err))
    return (None)

  for line in out.splitlines():
    if re.findall(u'Name:(.*)', line):
      vib_name = re.findall(u'Name:(.*)', line)[0]
      break
  if not vib_name:
    logging.error("Unable to get vib name from its vib file. out:%s" % out)
    return None
  return vib_name.strip()

def is_vib_installed(vib_name):
  """
  Check if VIB is currently installed in ESXi hypervisor.
  """
  cmd = "esxcli software vib get -n %s" % vib_name
  logging.info("Executing cmd: %s" % cmd)
  ret, out, err = run_command(cmd)
  if err or "[NoMatchError]" in out:
    logging.info("%s vib is not installed on Esxi host, ret %s out %s err %s"
                 % (vib_name, ret, out, err))
    return False
  logging.info("cmd %s, ret %s, out %s, err %s" % (cmd, ret, out, err))
  return True

def update_nutanix_vibs():
  """
  Update Nutanix VIBS which are installed during imaging.
  Returns status, reboot_required boolean flag upon installation of Vibs.
  """
  vib_files = glob.glob("%s/*.vib" % VIBS_DIR)
  success = True
  reboot_required = False
  for vib_file in vib_files:
    vib_name = get_vib_name(vib_file)
    if vib_name:
      need_to_install = not is_vib_installed(vib_name)
    else:
      logging.info("It could not fetch vib name from %s. Installing the vib." % vib_file)
      need_to_install = True
    if need_to_install:
      cmd = "esxcli software vib install -v=%s" % vib_file
    else:
      cmd = "esxcli software vib update -v=%s" % vib_file

    logging.info("Executing cmd: %s" % cmd )
    ret, out, err = run_command(cmd)
    if ret:
      logging.error("Unable to %s %s from cmd:%s,\nret: %s, out: %s, Error: %s"
                    % ("install" if need_to_install else "update",
                      vib_name, cmd, ret, out, err))
      success = False
    else:
      logging.info("VIB:%s is %s.\nOut:%s\nError:%s"
                   % (vib_name, "installed" if need_to_install else "updated",
                      out, err))
      pattern = re.compile("Reboot Required: true")
      match = pattern.search(out)
      if match:
        reboot_required = True
  return success, reboot_required

def modify_nutanix_boot_cfg(host):
  """
  Create a nutanix.tgz file and modify boot.cfg to invoke it.
  """
  tmp_dir = "/tmp"

  if host.summary.config.product.version < "5.1.0":
    rc_file = "/etc/rc.local"
  else:
    return True

  # Make /etc/rc.local.d/ directory path.
  path = os.path.join(tmp_dir, "etc/rc.local.d")
  if not os.path.exists(path):
    os.makedirs(path)

  rc_new = os.path.join(path, "nutanix.sh")
  shutil.copy(rc_file, rc_new)
  logging.info("Copied rc file to tmp dir")

  os.chdir(tmp_dir)
  cmd = "tar czpf nutanix.tgz etc"
  logging.info("Executing cmd: %s" % cmd )
  ret, out, err = run_command(cmd)
  if ret:
    logging.info("Creating tar file failed, ret %s, out %s, err %s" % (
                 ret, out, err))
    return False

  nutanix_zip_file = os.path.join(tmp_dir, "nutanix.tgz")
  esx_boot_dirs = [ "/bootbank", "/altbootbank" ]
  # Copy the nutanix.tgz file to /bootbank and /altbootbank.
  for boot_dir in esx_boot_dirs:
    shutil.copy(nutanix_zip_file, boot_dir)

  for boot_dir in esx_boot_dirs:
    boot_file = os.path.join(boot_dir, "boot.cfg")
    boot_cfg = None
    with open(boot_file, 'r') as fd:
      boot_cfg = fd.readlines()

    with NamedTemporaryFile(mode='w') as fd:
      for idx, line in enumerate(boot_cfg):
        if "modules=" in line:
          break

      line = boot_cfg[idx]
      if "nutanix.tgz" not in line:
        line = line.strip() + " --- nutanix.tgz\n"
      boot_cfg[idx] = line

      fd.seek(0, os.SEEK_SET)
      fd.truncate()
      fd.write("".join(boot_cfg))
      fd.flush()
      try:
        shutil.copy(fd.name, boot_file)
        logging.info("Copied boot cfg to %s" % boot_dir)
      except Exception as e:
        logging.info("Exception raised while copying boot cfg file %s" % e)
        return False

  shutil.rmtree(os.path.join(tmp_dir, "etc"))
  return True

def handle_vmd_reset_method(host):
  """
  Comment out entry of VMD reset method if upgrading to ESXi 7.0.
  Reboot of server is required when reset method is adjusted.
  Args:
      host: Requires to run PCI device discovery and to find ESXi
            version.
  Returns:
      bool: True if passthru.map is updated due to VMD else False.

  """
  def comment_purley_vmd(m):
    return '# Commented During Upgrade # ' + m.group(0)

  def is_vmd_present(host):
    hardware = host.hardware
    for pci in hardware.pciDevice:
      if ((pci.vendorId & 0xffff) == 0x8086 and
          (pci.deviceId & 0xffff) == 0x201d):
        return True
    return False

  if not is_vmd_present(host):
    logging.info("Purley VMD is not found. "
                 "No VMD reset method change required.")
    return False

  # Don't change passthru.map for ESXi 6.X
  if host.summary.config.product.version < "7.0.0":
    logging.info("ESXi version is lower than 7.0.0. "
                 "No VMD reset method change required.")
    return False

  # If there is "purley" VMD, comment it.
  with open("/etc/vmware/passthru.map") as pt_map:
    lines = pt_map.read()
  new_lines, fix_cnt = re.subn(r"^([ \t]*?8086\s+201d\s)", comment_purley_vmd,
                               lines, count=0, flags=re.M)
  if fix_cnt == 0:
    logging.info("No change in passthru.map required for VMD.")
    return False

  try:
    with NamedTemporaryFile(delete=False, mode='w') as fd:
      fd.write(new_lines)
      fd.flush()
      shutil.move(fd.name, "/etc/vmware/passthru.map")
    logging.info("Successfully updated reset method for VMD.")
    return True
  except shutil.Error as ee:
    logging.error("Shutil error seen while updating reset method for VMD: %s"
                  % ee)
    raise
  except (IOError, OSError) as ee:
    logging.error("IO error seen while updating reset method for VMD: %s"
                  % ee)
    raise
  return False

def modify_link_reset_method():
  """
  Modify the link reset method of PCI devices to bridge.
  """
  with open("/etc/vmware/passthru.map") as fd:
    data = fd.read()
  new_data = re.sub("1000  0097  (\S*)", "1000  0097  bridge", data)
  try:
    with NamedTemporaryFile(delete=False, mode='w') as fd:
      fd.write(new_data)
      fd.flush()
      shutil.move(fd.name, "/etc/vmware/passthru.map")
    logging.info("Successfully set the reset method to bridge for HBA")
  except shutil.Error as ee:
    logging.error("Shutil error seen during setting reset method to bridge %s"
                  % ee)
  except (IOError, OSError) as ee:
    logging.error("IO error seen during setting reset method to bridge %s"
                  % ee)
  return

def modify_sshd_config():
  """
  Modify /etc/ssh/sshd_config to remove "bad" MACs. ENG-155810.
  """
  sshd_config = "/etc/ssh/sshd_config"

  with open(sshd_config) as fd:
    data = fd.read()
  new_data = data.replace("hmac-sha1-96", "hmac-sha2-256")
  try:
    with NamedTemporaryFile(delete=False, mode='w') as fd:
      fd.write(new_data)
      fd.flush()
      shutil.move(fd.name, sshd_config)
    logging.info("Removed %s from %s" % ("hmac-sha1-96", "sshd_config"))
  except shutil.Error as ee:
    logging.error("Shutil error while changing sshd_config: %s" % ee)
  except (IOError, OSError) as ee:
    logging.error("IO error while changing sshd_config: %s" % ee)
  return

def vmd_marker_backup_file(target_dir):
  """
  Owner: Platforms Team
  Return path where VMD marker file should be backed up.
  """
  file_name = ""
  if not target_dir:
    logging.error("Cannot find svm_dir")
    return file_name

  just_file = os.path.basename(VMD_MARKER)
  file_name = os.path.join(target_dir, just_file)
  logging.info("VMD backup file is %s" % file_name)
  return file_name

def try_save_vmd_marker(target_dir):
  """
  Owner: Platforms Team
  VMD makrer is set only when CVM uses VMD. Since ESXi 7 upgrade does not
  keep /scratch partitions, need to preserve these marker by moving that
  in to local datastore, next to SVM files.
  Returns:
      True: Either VMD marker does not need to be managed or managing of
            marker is successful. Ongoing operation does not need to halt.
     False: Expected VMD marker to be backed up and failed in the operation.
            Ongoing operation should be stopped to avoid any potential
            regression.
  """
  retval = True
  if os.path.exists(VMD_MARKER):
    retval = False
    logging.info("VMD marker exist.")
    backup = vmd_marker_backup_file(target_dir)
    if not backup:
      logging.error("Unable to find valid backup path to save VMD marker "
                    "file.")
      return retval

    try:
      shutil.copy(VMD_MARKER, backup)
      logging.info("Saved %s file at %s" % (VMD_MARKER,
                                            backup))
      retval = True
    except Exception as e:
      logging.error("Exception raised while copying file %s, error: "
                    "%s" % (VMD_MARKER, e))
  else:
    logging.info("VMD marker does not exist.")

  return retval


def save_esx_upgrade_script(target_file):
  """
  Copy the esx_upgrade script to the cvm datastore just in case this happens to
  perform an upgrade to vSphere 7.0. vSphere 7.0 relays the /scratch partition,
  and as a result, all files created in /scratch partition will be lost after
  upgrade.
  """
  retval = False

  if not target_file:
    logging.error("Unable to find datastore location for CVM files")
    return retval

  try:
    shutil.copy(SOURCE_ESX_UPGRADE_FILE, target_file)
    logging.info("Saved %s file at %s" % (SOURCE_ESX_UPGRADE_FILE,
                                          target_file))
    retval = True
  except Exception as e:
    logging.error("Exception raised while copying file %s, error: "
                  "%s" % (SOURCE_ESX_UPGRADE_FILE, e))
  return retval

def save_file_to_datastore(source_file, target_dir):
  """
  Copy the source file to cvm datastore just in case this happens to
  perform an upgrade to vSphere 7.0. vSphere 7.0 relays the /bootbank and /scratch
  partition, and as a result, all files created in /bootbank  and /scratch partition
  will be lost after upgrade.
  """
  retval = False

  if not target_dir:
    logging.error("Unable to find datastore location for CVM files")
    return retval

  try:
    shutil.copy(source_file, target_dir)
    logging.info("Saved %s file at %s" % (source_file, target_dir))
    retval = True
  except Exception as e:
    logging.error("Exception raised while copying file %s to directory location %s, "
                  "error: %s" % (source_file, target_dir, e))
  return retval

def check_vmw_ahci_enabled():
  """
  check if vmw_ahci driver is enabled
  Return : (True, True) If driver is enabled and command
           executed successfully.
  """
  cmd = "esxcli system module list --enabled=true"
  ret, out, err = run_command(cmd)

  if ret:
    logging.error("Unable to run cmd: %s, ret: %s out: %s err: %s" % (cmd, ret,
                  out, err))
    return False, False

  logging.info("Checking for vmw_ahci driver enablement")
  for line in out.splitlines():
    if "vmw_ahci" in line:
      logging.info("vmw_ahci driver is enabled when performing esx upgrades")
      return True, True

  logging.info("vmw_ahci driver is not enabled when performing esx upgrades")
  return False, True


def enable_drivers():
  """
  Enable VMW AHCI driver if ESXI version greater than equal to 7.0
  """
  reboot_required = False
  os.remove(ENABLE_DRIVERS)
  vmw_ahci_enabled, ret = check_vmw_ahci_enabled()
  if not ret:
    logging.error("Failed to determine if vmw_ahci driver is "
                  "enabled when performing esx upgrades")
    sys.exit(0)

  if not vmw_ahci_enabled:
    cmd = "esxcli system module set --enabled=true --module=vmw_ahci"
    ret, out, err = run_command(cmd)
    if ret:
      logging.error("Unable to enable vmw_ahci driver, ret: "
                    "%d out: %s, err: %s" % (ret, out, err))
      sys.exit(0)
    else:
      logging.info("vmw_ahci driver is enabled successfully")
      reboot_required = True

  else:
    logging.info("vmw_ahci driver is already enabled when "
                 "performing esx upgrades")

  return reboot_required

def remove_vibs_pre_upgrade(host, target_version):
  """
  Remove vibs before upgrade.
  Args:
    host: Host object to get running esx version.
    target_version: Esx target version
  Return : (True, reboot_required) if vibs successfully uninstalled
  """
  # 1: Get vibs to remove
  vibs_to_remove = []
  if is_upgrading_to_esxi8_above(host.summary.config.product.version, target_version):
    vibs_to_remove.append(NFS_VAAI_VIB)
    vibs_to_remove.append(INTEL_OPTANE_VIB)
    vibs_to_remove.append(DEL_PTAGENT_VIB)
    vibs_to_remove.append(DEL_MNV_CLI_VIB)

  return remove_vibs(vibs_to_remove)

def remove_vibs(vibs_to_remove=None):
  """
  Remove vibs if installed
  vibs_to_remove: List of vibs to uninstall.
      If None we fetch deprecated vibs list form DEPRECATED_VIBS_LIST_PATH
  Return : (True, reboot_required) if vibs successfully uninstalled
          (False, reboot_required) if any of the vibs is failed to remove.
  """
  # 1: Get vibs to remove
  if vibs_to_remove is None:
    vibs_to_remove = []
    if os.path.exists(DEPRECATED_VIBS_LIST_PATH):
      try:
        f = open(DEPRECATED_VIBS_LIST_PATH)
        vibs_to_remove = json.load(f)
      except Exception as ee:
        logging.error('Unable to read %s file to get deprecated vibs list. Exception:%s'
                      % (DEPRECATED_VIBS_LIST_PATH, ee))
        return False, False

  # 2. Remove vibs if installed
  success = True
  reboot_required = False

  for vib in vibs_to_remove:
    cmd = REMOVE_VIB_CMD % vib
    logging.info("Executing cmd: %s" % cmd)
    (ret, out, err) = run_command(cmd)
    if not ret:
      logging.info("Successfully removed vib %s" % vib)
      pattern = re.compile("Reboot Required: true")
      match = pattern.search(out)
      if match:
        reboot_required = True
    else:
      no_match_error = 'NoMatchError'
      if no_match_error in out:
        logging.info("VIB %s not installed on the system, nothing to do" % vib)
      else:
        success = False
        logging.error("Failed to remove vib %s, ret: %s, out: %s, err: %s" %
                      (vib, ret, out, err))

  return success, reboot_required

def get_bundle_version_info():
  """
  Returns the version and build number (eg. 8.0.1) of the ESX host
  upgrade bundle located at \scratch\image.zip
  Return Type : (version, None) or (None, err_msg)
  """
  # Execute command on host console to find out build number.
  cmd = ("esxcli software sources vib get --depot=%s" % FILE)
  rv, out, err = run_command(cmd)
  if rv != 0:
    err = ("Failed to run cmd %r on host, ret %d, stdout %s, stderr %s" %
           (cmd, rv, out, err))
    logging.error(err)
    return (None, err)

  def cmp_key(xx):
    """
    If xx is ('8.0.1', '0.0', '21495797') for VMware_bootbank_esx-base_8.0.1-0.0.21495797
    return key would be [8, 0, 1, 0, 0, 21495797] for comparison
    Args:
      xx: version strings (tuple)
    Returns:
      numbers list for comparison
    """
    xx_list = []
    for ii in xx:
      xx_list = xx_list + ii.split(".")
    return ([int(x) for x in xx_list], None)

  # If there are multiple matching VIBs, the highest build number is valid.
  versions_list = re.findall(u'VMware_bootbank_esx-base_(.*)-(.*)\.(.*)', out)
  versions_list = sorted(versions_list, key=cmp_key)
  version, _, build_no = versions_list[-1]
  logging.info("esx bundle version is %s" % (version + '-' + build_no))
  return (version, None)

def is_upgrading_to_esxi8_above(running_version, target_version):
  """
  If host is getting upgraded to esxi 8.0.0 or above from lower than esxi 8.0.0
  """
  if (target_version and target_version >= "8.0.0" and
      running_version and running_version < "8.0.0"):
    return True
  return False

def upgrade(host, delay, vmlist, target_dir, target_file, profile_name):
  """
  Runs the upgrade command after delay. Before that we make sure that we are in
  maintenance mode. Delay default value is 30 secs. This is passed as an
  arguement from CVM.
  """
  try:
    time_delay = int(delay)
  except TypeError as ee:
    err_msg = "Could not parse delay with value '%s'" % delay
    logging.critical("%s Err: %s" % (err_msg, ee))
    mark_upgrade_failed(err_msg)
    sys.exit(1)

  error = True
  time.sleep(time_delay)
  timeout = MAINTENANCE_MODE_TIMEOUT
  # Timeout for host to go in maintenance mode.
  deadline = time.time() + timeout
  while time.time() < deadline:
    val = in_maintenance_mode(host)
    if not val:
      time.sleep(10)
      continue

    if not try_save_vmd_marker(target_dir):
      break

    if not save_esx_upgrade_script(target_file):
      break

    # Remove some VIBs prior to upgrade else upgrade might fail
    # Skipped the dry run for these error.
    # ENG-515005 , ENG-513311
    target_version, _ = get_bundle_version_info()
    if not os.path.exists(VIB_UPDATE_PENDING):
      remove_vibs_pre_upgrade(host, target_version)

    no_hardware_warning = ''
    if ((target_version and "8.0.1" <= host.summary.config.product.version < "9.0.0" and
        "8.0.2" <= target_version < "9.0.0" and os.path.exists(SKIP_HARDWARE_WARNING_FOR_SKYLAKE))
        or os.path.exists(SKIP_HARDWARE_WARNING_PATH)):
      logging.info("Skipping the hardware warnings")
      no_hardware_warning = " --no-hardware-warning"
    cmd = "esxcli software profile update -d=%s -p %s%s"\
          % (FILE, profile_name, no_hardware_warning)
    logging.info("Executing cmd: %s" % cmd)
    ret, out, err = run_command(cmd)
    if ret == 0:
      os.remove(SKIP_HARDWARE_WARNING_FOR_SKYLAKE) \
        if os.path.exists(SKIP_HARDWARE_WARNING_FOR_SKYLAKE) else None
      os.remove(SKIP_HARDWARE_WARNING_PATH) \
        if os.path.exists(SKIP_HARDWARE_WARNING_PATH) else None
      logging.info("Vib installation succeeded, Out:\n%s\nError:%s" % (
                   out, err))
      open(VIB_INSTALLED, 'w').close()
      pattern = re.compile("Reboot Required: true")
      match = pattern.search(out)
      if match and target_version and \
        is_upgrading_to_esxi8_above(host.summary.config.product.version, target_version):
        open(UPGRADE_MARKER, 'w').close()
        open(VIB_UPDATE_PENDING, 'w').close()
        with OpId():
          task = host.Reboot(force=True)
        # Due to vmware issue we need to reboot the host before updating the vibs
        # https://kb.vmware.com/s/article/90157
        logging.info("After host update reboot is required. "
                     "Rebooting host ret %s" % task.info.state)
        return
      # Remove installed deprecated VIBs before reboot.
      ret, reboot_vib_remove = remove_vibs()
      if not ret:
        err_msg = "Unable to remove installed VIBs prior to reboot"
        logging.error(err_msg)

      # Update VIB which are installed during imaging by Nutanix.
      ret, reboot_required_vibs = update_nutanix_vibs()
      os.remove(VIB_UPDATE_PENDING) if os.path.exists(VIB_UPDATE_PENDING) else None
      if match or reboot_required_vibs or reboot_vib_remove:
        logging.info("Reboot required")
        open(UPGRADE_MARKER, 'w').close()
        open(CLEANUP_PENDING, 'w').close()
        open(NIC_CHECK_MARKER, 'w').close()

        # Modify link reset method to Bridge.
        modify_link_reset_method()

        # Modify sshd_config to remove bad MACs.
        modify_sshd_config()

        # Make entry for nutanix.tgz in boot.cfg
        if not modify_nutanix_boot_cfg(host):
          # Creating nutanix.tgz failed
          err_msg = "Failed to update nutanix.tgz"
          logging.error(err_msg)
          post_upgrade_ops(host, target_dir, target_file)
          break
        else:
          cleanup_upgrade_pending()
          with OpId():
            task = host.Reboot(force=True)
          logging.info("Rebooting host. Task state: '%s'" % task.info.state)
          return
      else:
        # Reboot not required after VIB installation
        error = False
        post_upgrade_ops(host, target_dir, target_file)
        break
    else:
      # Vib installation failed
      error = True
      err_msg = "Vib installation failed"
      logging.error("%s, Return:%s Out:\n%s\nError:%s" % (err_msg, ret, out, err))
      post_upgrade_ops(host, target_dir, target_file)
      break

  # We could not enter maintenance mode.
  if not val:
    err_msg = "Timeout occurred while entering into maintenance mode"
    logging.error(err_msg)

  # Error occured during host upgrade.
  if error:
    mark_upgrade_failed(err_msg)

  # Remove the saved esx upgrade script
  remove_saved_esx_upgrade_script(target_file)
  try_remove_saved_vmd_marker(target_dir)

  cvm = nutanix_cvm(vmlist)
  if cvm:
    logging.info("Starting Nutanix CVM")
    task = cvm.PowerOn()
    logging.info("CVM poweron task state: %s" % task.info.state)
  else:
    logging.error("Cannot find nutanix CVM")
  return

def reboot_pending(host):
  """
  Return True if reboot is pending, else returns False. If reboot is pending,
  dont trigger upgrade.
  """
  return host.summary.rebootRequired

def reboot_host(host):
  """
  Reboot ESX host after waiting for maintenance mode task to complete.

  Args:
    host: Host object to be rebooted.

  Returns:
    True if the reboot operation was successfully started. A False is
    returned on errors.
  """
  timeout = MAINTENANCE_MODE_TIMEOUT
  # Timeout for host to go in maintenance mode.
  deadline = time.time() + timeout
  while time.time() < deadline:
    val = in_maintenance_mode(host)
    if not val:
      time.sleep(10)
      continue
    logging.info("Host in maintenance mode, Rebooting the host")
    if os.path.exists(REBOOT_MARKER):
      os.remove(REBOOT_MARKER)
    with OpId():
      task = host.Reboot(force=True)
    logging.info("Rebooting host. Task state: %s" % task.info.state)
    return True
  logging.error("Host not in maintenance mode after 10 minutes exiting")
  return False

def get_current_vibs():
  """
  Function that returns list of dictionary of VIB information.
  This function gurantees to have following format:
  <status(bool)>, [List_Of_Vibs_Name]
  """
  status = False
  command = "localcli software vib list"
  ret, out, err = run_command(command)
  if ret != 0:
    msg = "%s command failed.\nSTDOUT:\n%s\nSTDERR:\n\s" % (out, err)
    logging.error(msg)
    return status, []
  vib_list = []
  for line in out.splitlines()[2:]:
    information = line.split(" ", 1)
    vib_list.append(information[0])
  status = True
  return status, vib_list


def uninstall_vibs(list_of_vibs, from_altbank=False):
  """
  Uninstall list of vibs.
  Args:
      list_of_vibs: list of vibs to be removed.
      from_altbank(bool, False): Default value is false, remove from altbank.
                                 This strategy may save extra reboot.
  Returns:
      success, reboot_required
      success(bool): True if all vibs uninstalled else False
      reboot_required(bool): True if reboot is required else False
  """
  if not list_of_vibs:
    logging.info("No vibs to be removed.")
    return True, False
  logging.info("Followng vibs are to be uninstalled: %s" % list_of_vibs)

  if not from_altbank:
    remove_cmd = "localcli software vib remove"
  else:
    remove_cmd = "localcli software vib remove --no-live-install"
  command = ' -n '.join([remove_cmd] + list_of_vibs)

  reboot_required = False
  passed = True
  reboot_req_pattern = re.compile("Reboot Required: true")
  logging.info("Running %s" % command)
  ret, out, err = run_command(command)
  if ret == 0:
    logging.info("vib removal succeeded")
    match = reboot_req_pattern.search(out)
    if match:
      reboot_required = True
  else:
    logging.error("%s vib(s) removal failed, Out:\n%s\nError:%s" % (
      list_of_vibs, out, err)
    )
    passed = False

  return passed, reboot_required



def manage_vmd_vibs(host):
  """
  Note: This code is managed by platforms team.

  Expected to be called upon first boot after upgrade.
  This script shall attempt to take appropriate action
  to manage VMD vib(s).

  Args:
      host: Requires to get host version
  Returns:
    (status, reboot_required)
    status (bool): actions are successful when True
    reboot_required (bool): Reboot is required after this update when True.
    False: Reboot is not required after this update
  """
  if host.summary.config.product.version >= "8.0.0":
    logging.info("No vibs related to VMD reqiures uninstallation from esxi 8.0")
    return (True, False)

  vibs_to_remove_re = [r"iavmd", r"lsu.*-vmd-", r"intel-nvme-vmd*"]

  status, found_vibs = get_current_vibs()
  if not status:
    return False, False
  # Now run regex to find what VIBs are really required to be removed.
  vibs_to_remove = []
  for vib in found_vibs:
    for vtr in vibs_to_remove_re:
      if re.search(vtr, vib):
        vibs_to_remove.append(vib)
        break

  if not vibs_to_remove:
    logging.info("No vibs related to VMD reqiures uninstallation.")
    return (True, False)
  logging.info("Following vibs related to VMD are being "
               "uninstalled. %s" % vibs_to_remove)

  passed, reboot_required = uninstall_vibs(vibs_to_remove)

  logging.info("VMD VIB Management: passed=%s, reboot_required=%s" % (
    passed,
    reboot_required
  ))

  return (passed, reboot_required)

def get_target_directory_and_file(vmlist):
  """
  Method to get target directory and target file.
  Also, it takes backup of config files. 
  Args:
    vmlist: List of VMs
  Returns:
    Target directory and file for CVM datastore.
  """
  try:
    cvm = nutanix_cvm(vmlist)
    # Derive its datastore
    if cvm is not None:
      cvm_ds = cvm.config.datastoreUrl[0].name
      target_dir = get_svm_dir(cvm_ds)
    else:
      # Save the config files to VMFS datastore if CVM is
      # not found.
      uuid = open(CO_DATASTORE_UUID, "r").read().splitlines()[0]
      target_dir = get_svm_dir_co(uuid)
      save_file_to_datastore(CO_DATASTORE_UUID, target_dir)
      save_file_to_datastore(SOURCE_HARDWARE_CONFIG_FILE, target_dir)
      save_file_to_datastore(SOURCE_FACTORY_CONFIG_FILE, target_dir)
    logging.info("Target directory is: %s" % target_dir)
    target_file = os.path.join(target_dir, '.esx_upgrade')
    logging.info("Target file is: %s" % target_file)
    return (target_dir, target_file)
  except Exception as e:
    err_msg = "Unable to determine CVM datastore name, exception %s" % e
    logging.error(err_msg)
    mark_upgrade_failed(err_msg)
    sys.exit(1)

def main():
  if len(sys.argv) < 2:
    err_msg = "Incorrect invocation of %s" % sys.argv[0]
    print(err_msg)
    logging.info(err_msg)
    mark_upgrade_failed(err_msg)
    sys.exit(1)
  delay = sys.argv[1]
  profile_name = ""
  if len(sys.argv) > 2:
    profile_name = sys.argv[2]

  logging.info("Delay: '%s', profile: '%s'" % (delay, profile_name))

  # Adds support for using ssl in python 3.X
  try:
    original_match_hostname = ssl.match_hostname

    @functools.wraps(original_match_hostname)
    def skip_match_hostname_check(cert, hostname):
      return True

    ssl.match_hostname = skip_match_hostname_check
  except AttributeError:
    pass

  # Python 3.X we need to update the sslContext, but that is not possible in
  # Python 2.X
  try:
    stub = SoapAdapter.SoapStubAdapter(host="localhost", port=443,
        version="vim.version.version7", path="/sdk",
        sslContext=ssl._create_unverified_context())
  except (TypeError, AttributeError):
    stub = SoapAdapter.SoapStubAdapter(host="localhost", port=443,
        version="vim.version.version7", path="/sdk")

  output = {}
  deadline = time.time() + 900
  target_dir = ""
  host = None
  vmlist = None
  while time.time() < deadline:
    try:
      content = vim.ServiceInstance("ServiceInstance", stub).RetrieveContent()
      ticket = content.sessionManager.AcquireLocalTicket(os.environ["USER"])

      output = { "user": ticket.userName,
                 "pwd": open(ticket.passwordFilePath).read() }
    except HTTPException as e:
      logging.error("Http Exception occured %s" % e)
      time.sleep(1)
      continue
    
    if not output:
      logging.error("Unable to connect to vim service instance")
      time.sleep(1)
      continue
    
    try:
      with OpId(): 
        si = SmartConnect(user=output["user"], pwd=output["pwd"])
    except vim.fault.HostConnectFault as ee:
      logging.error("Connection to host failed err %s" % ee.msg)
      time.sleep(1)
      continue

    # vSphere 7 changes the storage partition layout. Backup config files
    # to handle this.
    try:
      sc = si.content
      dc = sc.rootFolder.childEntity[0]
      cr = dc.hostFolder.childEntity[0]
      host = cr.host[0]
      vmlist = dc.vmFolder.childEntity
      break
    except Exception as e:
      logging.error("Unable to fetch the VM list, exception %s" % e)
      time.sleep(10)
  else:
    err_msg = "Unable to fetch the VM list"
    logging.error(err_msg)
    mark_upgrade_failed(err_msg)
    sys.exit(1)

  if reboot_pending(host):
    logging.info("Reboot required is set to true, postponing upgrade"
                 " till reboot")
    open(UPGRADE_MARKER, 'w').close()
    if not reboot_host(host):
      err_msg = "[reboot_pending] Host did not enter into maintenance mode"
      mark_upgrade_failed(err_msg)
      sys.exit(1)
    return

  if os.path.exists(ENABLE_DRIVERS):
    reboot_required = enable_drivers()
    if reboot_required:
      open(UPGRADE_MARKER, 'w').close()
      if not reboot_host(host):
        err_msg = "[enable_drivers] Host did not enter into maintenance mode"
        mark_upgrade_failed(err_msg)
        sys.exit(1)
      return

  if os.path.exists(REBOOT_MARKER):
    # Wait for host to go to maintenance mode, and reboot. The marker is
    # removed in reboot function only in positive scenario.
    if not reboot_host(host):
      err_msg = "[reboot_marker] Host did not enter into maintenance mode"
      mark_upgrade_failed(err_msg)
      sys.exit(1)
    os.remove(REBOOT_MARKER) if os.path.exists(REBOOT_MARKER) else None
  elif os.path.exists(UPGRADE_PENDING) or os.path.exists(VIB_UPDATE_PENDING):
    target_dir, target_file = get_target_directory_and_file(vmlist)
    if not profile_name:
      err_msg = "Profile name not provided for upgrade"
      logging.error(err_msg)
      mark_upgrade_failed(err_msg)
      sys.exit(1)
    upgrade(host, delay, vmlist, target_dir, target_file, profile_name)
    cleanup_upgrade_pending()
  elif os.path.exists(CLEANUP_PENDING):
    # This section of code is potentially recursive in nature.
    # Normal theme is - whenever reboot_required, one or more funcitons
    # performed their work and identified that reboot is required
    # Post reboot, same function shall be called and it is expected
    # that function return "pass" _and_ reboot due to that function call
    # be false. Having said, that if that function itself couldn't it its
    # objective, upgrade will fail.
    target_dir, target_file = get_target_directory_and_file(vmlist)
    reboot_required = False
    markers_to_remove_upon_reboot = []
    # Check for NIC connectivity.
    if os.path.exists(NIC_CHECK_MARKER):
      logging.info("Check NIC connectivity, and whether ixgben driver needs "
                   "to be enabled")

      connectivity_pass, reboot_due_to_nic = check_nic_connectivity()
      reboot_required = reboot_required or reboot_due_to_nic

      if not connectivity_pass:
        err_msg = "Network connectivity check failed."
        logging.error(err_msg)
        os.remove(NIC_CHECK_MARKER)
        mark_upgrade_failed(err_msg)
        sys.exit(1)
      markers_to_remove_upon_reboot.append(NIC_CHECK_MARKER)

    if os.path.exists(VMD_MARKER):
      logging.info("Check if VMD vibs need to be managed")
      vmd_vibs_success, reboot_due_to_vmd = manage_vmd_vibs(host)
      reboot_required = reboot_required or reboot_due_to_vmd

      pt_map_updated = handle_vmd_reset_method(host)
      reboot_required = reboot_required or pt_map_updated

      if not vmd_vibs_success:
        err_msg = "Error managing VMD vibs"
        logging.error(err_msg)
        os.remove(VMD_MARKER)
        mark_upgrade_failed(err_msg)
        sys.exit(1)
      markers_to_remove_upon_reboot.append(VMD_MARKER)

    for marker in markers_to_remove_upon_reboot:
      if os.path.exists(marker):
        os.remove(marker)
    # Accumulate all reboots and issue only once.
    if reboot_required:
      # local.sh removes this marker unconditionally.
      # Not having this marker will not call this script again.
      open(UPGRADE_MARKER, 'w').close()
      with OpId():
        task = host.Reboot(force=True)
      logging.info("Rebooting host ret %s" % task.info.state)
      # Reboot action takes long time after issueing command.
      # It is not wise to start executing next set of instructions
      # when reboot is pending. Next "boot" will call this function
      # again.
      sys.exit(0)

    post_upgrade_ops(host, target_dir, target_file)
    os.remove(CLEANUP_PENDING)
    open(UPGRADE_FINISHED, 'w').close()
    logging.info("Created upgrade finished marker")

  elif os.path.exists(MAINTENANCE_CLEANUP_PENDING):
    maintenance_cleanup(host)
    os.remove(MAINTENANCE_CLEANUP_PENDING)
  else:
    print("Unknown arguments")

  if os.path.exists(PCI_RECONFIG_MARKER):
    logging.info("Running esx_pci_reconfig_cvm script")
    # This script can reboot host.
    ret, out, err = run_command("python /scratch/esx_pci_reconfig_cvm")
    if ret != 0:
      logging.error("Unable to run cvm pci reconfig, ret %s out %s err %s" % (
                    ret, out, err))
      return
    logging.info("Successfully run esx_pci_reconfig_cvm script")
  return

if __name__ == "__main__":
  try:
    logging.basicConfig(format="%(asctime)s %(levelname)s "
                               "%(funcName)s:%(lineno)d - %(message)s",
                        filename=LOGFILE,
                        level=logging.DEBUG)
    main()
  except Exception:
    err_msg = "Failed to run esx_upgrade"
    logging.exception(err_msg)
    mark_upgrade_failed(err_msg)
    sys.exit(1)

