########################################################################
# Copyright (c) 2018-2019 VMware, Inc.  All rights reserved
# VMware Confidential
########################################################################

"""Create, modify, and/or format UEFI boot options and boot order.
See "UEFI: Unified Extensible Firmware Interface Specification",
version 2.7, May 2017, chapter 3, "Boot Manager".

"""

import struct
import re
import pprint
import uuid
import ipaddress

from . import uefivar

if uefivar.isVMkernel:
   from vmware import vsi
   from vmware import runcommand

class DevicePathError(Exception):
   """Error parsing a UEFI device path.
   """
   pass

BOOT_ORDER = 'BootOrder-' + uefivar.EfiGlobalVariableGUID
BOOT_NEXT = 'BootNext-' + uefivar.EfiGlobalVariableGUID
BOOT_CURRENT = 'BootCurrent-' + uefivar.EfiGlobalVariableGUID
BOOT_OPTION_PAT = 'Boot([0-9A-Fa-f]{4})-' + uefivar.EfiGlobalVariableGUID

VAR_ATTRIBS = (uefivar.EfiVariableNonVolatile |
               uefivar.EfiVariableBootserviceAccess |
               uefivar.EfiVariableRuntimeAccess)

# EFI_LOAD_OPTION attributes (UEFI 2.7 section 3.1.3)
LOAD_OPTION_ACTIVE          = 0x00000001
LOAD_OPTION_FORCE_RECONNECT = 0x00000002
LOAD_OPTION_HIDDEN          = 0x00000008
LOAD_OPTION_CATEGORY_SHIFT  = 8
LOAD_OPTION_CATEGORY        = 0x00001F00
LOAD_OPTION_CATEGORY_BOOT   = 0x00000000
LOAD_OPTION_CATEGORY_APP    = 0x00000100

# Disk or partition signature
MBR_SIGNATURE_OFFSET = 0x1b8
MBR_SIGNATURE_LEN    = 4
GPT_SIGNATURE_LEN    = 16

VSI_SBDF_PATH = '/hardware/pci/seg/{seg}/bus/{bus}/slot/{dev}/func/{func}'

# Some GUIDs for formatDevPath
RAMDISK_GUIDs = {
   '77ab535a-45fc-624b-5560-f7b281d1f96e': 'VirtualDisk',
   '3d5abd30-4175-87ce-6d64-d2ade523c4bb': 'VirtualCD',
   '5cea02c9-4d07-69d3-269f-4496fbe096f9': 'PersistentVirtualDisk',
   '08018188-42cd-bb48-100f-5387d53ded3d': 'PersistentVirtualCD'
}

#
# Creating, parsing, and formatting UEFI device paths
#

def unpackUUID(data):
   """Convert 16 byte little-endian UUID to an uppercase string.
   """
   return str(uuid.UUID(bytes_le=data)).upper()

def unpackEisaId(eisaid):
   """Convert UEFI compressed EISA-type ID to a string.  See UEFI spec
   2.7 section 10.3.3.
   """
   return (chr(((eisaid >> 0) & 0x1f) + ord('A') - 1) +
           chr(((eisaid >> 5) & 0x1f) + ord('A') - 1) +
           chr(((eisaid >> 10) & 0x1f) + ord('A') - 1) +
           ('%04X' % (eisaid >> 16)))

def packEisaId(s):
   """Convert UEFI EISA-type ID string to an integer.  See UEFI spec 2.7
   section 10.3.3.
   """
   return (((ord(s[0]) - ord('A') + 1) << 0) +
           ((ord(s[1]) - ord('A') + 1) << 5) +
           ((ord(s[2]) - ord('A') + 1) << 10) +
           (int(s[3:], 16) << 16))

def packHDDevPath(partNum, gpt, sig, startLB, sizeLB):
   """Pack parameters into a hard drive media device path.  See UEFI Spec
   2.7 Table 86 "Hard Drive Media Device Path".
   """
   partFmt = 2 if gpt else 1
   sigType = (2 if len(sig) == GPT_SIGNATURE_LEN else
              1 if len(sig) == MBR_SIGNATURE_LEN else
              0)
   dp = struct.pack('<BBHIQQ16sBB',
                    4,       # Type 4 - Media Device Path
                    1,       # Sub-Type 1 - Hard Drive
                    42,      # Length of this structure in bytes.
                    partNum, # Partition Number
                    startLB, # Partition Start (LBA)
                    sizeLB,  # Partition Size (in LB)
                    sig,     # Partition Signature
                    partFmt, # Partition Format
                    sigType) # Signature Type
   return dp

def packStr(s):
   """Pack a string as utf-16-le, adding an explicit NUL terminator.
   """
   return (s + '\x00').encode('utf-16-le')

def unpackStr(b):
   """Unpack a utf-16-le string, stripping explicit NUL terminator if
   present.
   """
   s = b.decode('utf-16-le')
   if s[-1:] == '\x00':
      s = s[:-1]
   return s

def packFPDevPath(pathName):
   """Pack a pathname into a file path media device path.  See UEFI Spec
   2.7 Table 89 "File Path Media Device Path".
   """
   ps = packStr(pathName)
   ln = 4 + len(ps)
   dp = struct.pack('<BBH',
                    4,       # Type 4 - Media Device Path
                    4,       # Sub-Type 4 - File Path
                    ln)      # Length of this structure in bytes.
   return dp + ps

def packEndPath(instance=False):
   """Pack parameters into an end device path.  See UEFI Spec 2.7 Table
   41 "Device Path End Structure".
   """
   subTyp = 0x01 if instance else 0xff
   dp = struct.pack('<BBH',
                    0x7f,    # Type 0x7F - End of Hardware Device Path
                    subTyp,
                    4)       # Length of this structure in bytes.
   return dp

def packPCIDevPath(sbdf):
   """Pack a PCI device into a device path. See UEFI Spec 2.7 Table 48
   "ACPI Device Path" and UEFI Spec 2.7 Table 42 "PCI Device Path".

   Parameters: sbdf must be a string of the form 0000:00:00.0
   """
   # Split sbdf
   node = {}
   (node['seg'], node['bus'], node['dev'], node['func']) = \
      [int(x, 16) for x in re.split('[:.]', sbdf)]

   # Emit correct PciRoot for this device
   #
   # PNP0A03 = PCI Bus; see http://www.uefi.org/PNP_ACPI_Registry.
   # Technically PcieRoot (PNP0A08) would be more accurate for current
   # machines that use PCIe, but on most machines it seems that the
   # firmware uses PciRoot (PNP0A03).  Maybe they are interchangeable
   # in practice.
   #
   hid = packEisaId('PNP0A03')
   uid = vsi.get((VSI_SBDF_PATH + '/rootBridgeUID').format(**node))
   dp = struct.pack('<BBHII',
                    2,        # Type 2 - ACPI Device Path
                    1,        # Sub-Type 1 - ACPI Device Path
                    12,       # Length of this structure in bytes.
                    hid,      # EISA compressed PnP hardware ID
                    uid)      # Unique ID

   # Walk PCI hierarchy upward from sbdf
   pciPath = []
   while node['bus'] != (1 << 32) - 1:
      pciPath.append(node)
      node = vsi.get((VSI_SBDF_PATH + '/parent').format(**node))

   # Walk back down, emitting PCI(func,dev) for each node
   for node in reversed(pciPath):
      pr = struct.pack('<BBHBB',
                       1,        # Type 1 - Hardware Device Path
                       1,        # Sub-Type 1 - PCI
                       6,        # Length of this structure in bytes
                       node['func'],
                       node['dev'])
      dp = dp + pr

   return dp

def packMACDevPath(mac, macType=1):
   """Pack parameters into a MAC address device path.  See UEFI Spec 2.7
   Table 65 "MAC Address Device Path".

   Parameters:
   mac: must be a string of the form '12:34:56:78:9a:bc'
   macType: DHCP hw address type, 0 (unspecified) or 1 (Ethernet)
   """
   dp = (struct.pack('<BBH',
                     3,       # Type 3 - Messaging Device Path
                     11,      # Sub-Type 11 - MAC Address
                     37) +    # Length of this structure in bytes.
         bytes([int(b, 16) for b in mac.split(':')]) +
         b'\x00' * 26 +       # Padding
         bytes(macType))      # DHCP hardware address type
   return dp

def packIPv4DevPath():
   """Pack an empty IPv4 device path.  That is, this path just says "use
   IPv4"; it doesn't say what addresses to use, etc.
   """
   dp = (struct.pack('<BBH',
                     3,       # Type 3 - Messaging Device Path
                     12,      # Sub-Type 12 - IPv4
                     27) +    # Length of this structure in bytes.
         b'\x00' * 23)
   return dp

def packIPv6DevPath():
   """Pack an empty IPv6 device path.  That is, this path just says "use
   IPv6"; it doesn't say what addresses to use, etc.
   """
   dp = (struct.pack('<BBH',
                     3,       # Type 3 - Messaging Device Path
                     13,      # Sub-Type 13 - IPv6
                     60) +    # Length of this structure in bytes.
         b'\x00' * 56)
   return dp

def packURIDevPath(uri):
   """Pack a URI device path. The uri must be plain ASCII, with
   percent-encoding already applied if needed to represent reserved or
   non-ASCII characters.  See RFC3986 and
   https://en.wikipedia.org/wiki/Percent-encoding.
   """
   u = uri.encode('ascii')
   ln = 4 + len(u)
   dp = (struct.pack('<BBH',
                     3,       # Type 3 - Messaging Device Path
                     24,      # Sub-Type 24 - Universal Resource Identifier
                     ln) +    # Length of this structure in bytes.
         u)
   return dp

def bytesToHex(bb):
   """Convert bytes to an uppercase hex string.
   """
   return ''.join(['%02X' % b for b in bb])

def formatDevPath(dp):
   """Convert UEFI file/device path list to human-readable text.  For
   debug and display purposes only.  Implements UEFI 2.8 section 10.6,
   table 105, except that many specific node types are not implemented
   and are thus shown as a more generic supertype.  Please add more
   cases as needed.

   "End this instance of a device path" is shown as ',' and "End
   entire device path" is shown as ';'.
   """
   i = 0
   path = ''
   gotEnd = False
   while i < len(dp):
      (typ, subTyp, length) = struct.unpack('<BBH', dp[i : i + 4])
      data = dp[i + 4 : i + length]

      if typ == 1:      # Hardware Device Path
         if subTyp == 1:    # PCI
            node = 'Pci(0x%X,0x%X)' % (data[1], data[0])
         elif subTyp == 3:  # Memory Mapped
            node = 'MemoryMapped(0x%X,0x%X,0x%X)' % struct.unpack('<IQQ', data)
         elif subTyp == 4:  # Vendor
            node = 'VenHw(%s,%s)' % (unpackUUID(data[:16]),
                                     bytesToHex(data[16:]))
         elif subTyp == 5:  # Controller
            node = 'Ctrl(0x%X)' % struct.unpack('<I', data)
         else:
            node = 'HardwarePath(%d,%s)' % (subTyp, bytesToHex(data))

      elif typ == 2:    # ACPI Device Path
         if subTyp == 1:    # ACPI Device Path
            (hid, uid) = struct.unpack('<II', data)
            hid = unpackEisaId(hid)
            if hid == 'PNP0A03':
               node = 'PciRoot(0x%X)' % uid
            elif hid == 'PNP0A08':
               node = 'PcieRoot(0x%X)' % uid
            else:
               node = 'Acpi(%s,0x%X)' % (hid, uid)
         else:
            node = 'AcpiPath(%d,%s)' % (subTyp, bytesToHex(data))

      elif typ == 3:    # Messaging Device Path
         if subTyp == 1:    # ATAPI
            node = 'Ata(0x%X,0x%X,0x%X)' % struct.unpack('<BBH', data)
         elif subTyp == 2:  # SCSI
            node = 'Scsi(0x%X,0x%X)' % struct.unpack('<HH', data)
         elif subTyp == 3:  # Fibre Channel
            node = 'Fibre(0x%X,0x%X)' % struct.unpack('<IQQ', data)[1:3]
         elif subTyp == 21: # Fibre Channel Ex
            node = 'FibreEx(%s,%s)' % (bytesToHex(data[4:12]),
                                       bytesToHex(data[12:20]))
         elif subTyp == 5:  # USB
            node = 'USB(0x%X,0x%X)' % struct.unpack('<BB', data)

         elif subTyp == 11: # MAC Address
            ifType = data[-1]
            if ifType == 0 or ifType == 1:
               macLen = 6
            else:
               macLen = 32
            node = 'MAC(%s,0x%X)' % (bytesToHex(data[:macLen]), ifType)

         elif subTyp == 12: # IPv4
            if len(data) == 23:
               (lip, rip, lpt, rpt, proto, orig, gip, mip) = \
                  struct.unpack('<4s4sHHHB4s4s', data)
            else:
               (lip, rip, lpt, rpt, proto, orig) = \
                  struct.unpack('<4s4sHHHB', data)
               gip = b'\x00' * 4
               mip = b'\x00' * 4
            lpt = '' if lpt == 0 else ':' + str(lpt)
            rpt = '' if rpt == 0 else ':' + str(rpt)
            if proto == 17:
               proto = 'UDP'
            elif proto == 6:
               proto = 'TCP'
            else:
               proto = hex(proto)
            if orig <= 1:
               orig = ('DHCP', 'Static')[orig]
            else:
               orig = hex(orig)
            node = ('IPv4(%s%s,%s,%s,%s%s,%s,%s)' %
                    (str(ipaddress.IPv4Address(rip)), rpt, proto, orig,
                     str(ipaddress.IPv4Address(lip)), lpt,
                     str(ipaddress.IPv4Address(gip)),
                     str(ipaddress.IPv4Address(mip))))

         elif subTyp == 13: # IPv6
            if len(data) == 56:
               (lip, rip, lpt, rpt, proto, orig, prelen, gip) = \
                  struct.unpack('<16s16sHHHBB16s', data)
            else:
               (lip, rip, lpt, rpt, proto, orig) = \
                  struct.unpack('<16s16sHHHB', data)
               prelen = 0
               gip = b'\x00' * 16
            lpt = '' if lpt == 0 else ':' + str(lpt)
            rpt = '' if rpt == 0 else ':' + str(rpt)
            if proto == 17:
               proto = 'UDP'
            elif proto == 6:
               proto = 'TCP'
            else:
               proto = hex(proto)
            if orig <= 2:
               orig = ('Static', 'StatelessAutoConfigure',
                       'StatefulAutoConfigure')[orig]
            else:
               orig = hex(orig)
            node = ('IPv6(%s%s,%s,%s,%s%s,0x%x,%s)' %
                    (str(ipaddress.IPv6Address(rip)), rpt, proto, orig,
                     str(ipaddress.IPv6Address(lip)), lpt,
                     prelen, str(ipaddress.IPv6Address(gip))))

         elif subTyp == 17: # Device Logical Unit
            node = 'Unit(0x%X)' % data[0]
         elif subTyp == 18: # SATA
            node = 'Sata(0x%X,0x%X,0x%X)' % struct.unpack('<HHH', data)
         elif subTyp == 24: # URI
            # Using latin-1 here for paranoia; ASCII should be enough.
            # See RFC3986 and https://en.wikipedia.org/wiki/Percent-encoding
            node = 'Uri(%s)' % data.decode('latin-1')
         else:
            node = 'Msg(%d,%s)' % (subTyp, bytesToHex(data))

      elif typ == 4:    # Media Device Path
         if subTyp == 1:    # Hard Drive
            (partNum, start, size, sig, fmt, sigType) = \
               struct.unpack('<IQQ16sBB', data)
            fmtStr = ('MBR' if fmt == 1 else
                      'GPT' if fmt == 2 else
                      '0x%X' % fmt)
            sigStr = (unpackUUID(sig) if sigType == 2 else
                      '0x%X' % struct.unpack('<I', sig[:4]) if sigType == 1 else
                      '0' if sigType == 0 else
                      bytesToHex(sig))
            node = 'HD(0x%X,%s,%s,0x%X,0x%X)' % \
                   (partNum, fmtStr, sigStr, start, size)

         elif subTyp == 2:  # CD-ROM
            node = 'CDROM(0x%X,0x%X,0x%X)' % struct.unpack('<IQQ', data)
         elif subTyp == 4:  # String
            node = data.decode('utf-16-le')
         elif subTyp == 6:  # PIWG Firmware File
            node = 'FvFile(%s)' % unpackUUID(data)
         elif subTyp == 7:  # PIWG Firmware Volume
            node = 'Fv(%s)' % unpackUUID(data)
         elif subTyp == 8:  # Relative Offset Range
            (reserved, start, end) = struct.unpack('<IQQ', data)
            node = 'Offset(0x%X,0x%X)' % (start, end)
         elif subTyp == 9:  # RAM Disk
            (start, end, dtype, inst) = struct.unpack('<QQ16sH', data)
            u = unpackUUID(dtype)
            try:
               name = RAMDISK_GUIDs[u.lower()]
               node = '%s(0x%X,0x%X,0x%X)' % (name, start, end, inst)
            except KeyError:
               node = 'RamDisk(0x%X,0x%X,0x%X,%s)' % (start, end, inst, u)
         else:
            node = 'MediaPath(%d,%s)' % (subTyp, bytesToHex(data))

      elif typ == 5:    # BIOS Boot Specification Device Path
         if subTyp == 1:
            # Using latin-1 here for paranoia; UEFI spec says ASCII.
            (d, f) = struct.unpack('<HH', data[:4])
            node = 'BBS(0x%X,%s,0x%X)' % (d, data[4:].decode('latin-1'), f)
         else:
            node = 'BbsPath(%d,%s)' % (subTyp, bytesToHex(data))

      elif typ == 0x7f: # End of Hardware Device Path
         if subTyp == 1:      # End This Instance of a Device Path
            node = ','
         elif subTyp == 0xff: # End Entire Device Path
            node = ';'

      else:
         node = 'Path(%d,%d,%s)' % (typ, subTyp, bytesToHex(data))

      path = path + node
      i += length
      if length == 0 or i > len(dp):
         raise DevicePathError('Invalid length field in UEFI device path')

   return path

#
# Boot option support
#

def bootOptionName(optionNumber):
   """Return the name of the UEFI variable for the specified boot option.
   """
   return 'Boot%04X-%s' % (optionNumber, uefivar.EfiGlobalVariableGUID)

def getRawBootOption(optionNumber):
   """Get the raw value of the specified boot option.
   """
   return uefivar.get(bootOptionName(optionNumber))

def setRawBootOption(optionNumber, raw):
   """Set the raw value of the specified boot option.
   """
   uefivar.set(bootOptionName(optionNumber), raw)

def avoidHexAF(n):
   """Return the first integer m >= n such that m does not have any
   digits in the range [A-F] in its hex representation.  This function
   is used to avoid a bug in certain older UEFI implementations.
   """
   i = 0
   while n > (9 << i):
      if ((n >> i) & 0xf) > 9:
         n = (n & (-1 << (i + 4))) + (1 << (i + 4))
      i += 4
   return n

def getBootOrder():
   """Get the boot order as a list of integer option numbers.
   """
   bootOrder = uefivar.get(BOOT_ORDER)
   bootOrderIt = struct.iter_unpack('<H', bootOrder[4:])
   return [elt[0] for elt in bootOrderIt]

def setBootOrder(order):
   """Set a new boot order (using existing boot options).  The argument
   is a list of integer option numbers.
   """
   val = struct.pack('<I', VAR_ATTRIBS)
   for i in order:
      val += struct.pack('<H', i)
   uefivar.set(BOOT_ORDER, val)

def getAllBootOptions():
   """Return a list of all currently defined boot options, whether in the
   boot order or not, sorted numerically.
   """
   names = uefivar.list()
   nums = []
   for name in names:
      match = re.match(BOOT_OPTION_PAT, name)
      if match:
         nums.append(int(match.groups()[0], 16))
   return sorted(nums)

def getBootNext():
   """Get one-time boot option number, if any.  If there is none, return None.
   """
   try:
      raw = uefivar.get(BOOT_NEXT)
   except uefivar.UefiVarNotFound:
      return None
   return struct.unpack('<IH', raw)[1]

def setBootNext(optionNumber):
   """Set up one-time boot from the given option number.  If optionNumber
   is None, disable one-time boot.
   """
   if optionNumber is None:
      try:
         uefivar.remove(BOOT_NEXT)
      except uefivar.UefiVarNotFound:
         pass
   else:
      val = struct.pack('<IH', VAR_ATTRIBS, optionNumber)
      uefivar.set(BOOT_NEXT, val)

def promoteBootOption(optionNumber):
   """Ensure optionNumber is first in the boot order, adding it if
   needed.
   """
   bootOrder = getBootOrder()
   if optionNumber in bootOrder:
      bootOrder.remove(optionNumber)
   bootOrder.insert(0, optionNumber)
   setBootOrder(bootOrder)

def demoteBootOption(optionNumber):
   """Ensure optionNumber is last in the boot order, adding it if
   needed.
   """
   bootOrder = getBootOrder()
   if optionNumber in bootOrder:
      bootOrder.remove(optionNumber)
   bootOrder.append(optionNumber)
   setBootOrder(bootOrder)

def removeBootOption(optionNumber, clean=False):
   """Remove an option from the boot order if present.  If clean=True,
   remove the corresponding BootNNNN variable as well.
   """
   bootOrder = getBootOrder()
   if optionNumber in bootOrder:
      bootOrder.remove(optionNumber)
      setBootOrder(bootOrder)
   if clean:
      uefivar.remove(bootOptionName(optionNumber))

def getBootCurrent():
   """Get integer boot option number that we are currently booted from.
   """
   raw = uefivar.get(BOOT_CURRENT)
   return struct.unpack('<IH', raw)[1]

def packBootOption(option):
   """Pack a boot option dict into a raw value.

   Input:
      varAttr,     # Attributes of variable
      optionAttr,  # Attributes of boot option
      description  # Description string (decoded from UTF-16)
      filePathList # UEFI file path list
      optionalData # Raw optional data

   Output: bytes object (*including* the UEFI variable attributes)
   """
   filePathListLen = len(option['filePathList'])
   return (struct.pack('<IIH', option['varAttr'],
                       option['optionAttr'], filePathListLen) +
           packStr(option['description']) +
           option['filePathList'] +
           option['optionalData'])

def unpackBootOption(raw):
   """Unpack one boot option into its fields, returning a dict.

   Input: bytes object (*including* the UEFI variable attributes)

   Output:
      varAttr,     # Attributes of variable
      optionAttr,  # Attributes of boot option
      description  # Description string (decoded from UTF-16)
      filePathList # UEFI file path list
      optionalData # Raw optional data
   """
   (varAttr, optionAttr, filePathListLen) = struct.unpack('<IIH', raw[:10])
   match = re.search(b'(^(..)*?)\x00\x00', raw[10:])
   description = unpackStr(match.group(1))
   filePathList = raw[10 + match.end() : 10 + match.end() + filePathListLen]
   optionalData = raw[10 + match.end() + filePathListLen : ]
   return {'varAttr':      varAttr,    # Attributes of variable
           'optionAttr':   optionAttr, # Attributes of boot option
           'description':  description,
           'filePathList': filePathList,
           'optionalData': optionalData}

def createRawBootOption(raw):
   """Create a boot option with the given raw value, suppressing
   duplicates and returning the option number.  "Suppressing
   duplicates" means that if an option with the given raw value
   already exists, its number is simply returned.  Otherwise the
   option is created with a new, currently unused option number.

   To avoid a bug in certain older UEFI implementations, avoids creating
   a new option number that has a hex digit outside the [A-F] range.
   """
   newOptNum = 0x0000
   names = uefivar.list()
   for name in sorted(names):
      match = re.match(BOOT_OPTION_PAT, name)
      if match:
         optNum = int(match.groups()[0], 16)
         if uefivar.get(name) == raw:
            return optNum
         if newOptNum == optNum:
            newOptNum = avoidHexAF(optNum + 1)
   setRawBootOption(newOptNum, raw)
   return newOptNum

def createBootOption(description, filePathList,
                     optionAttr=LOAD_OPTION_ACTIVE, optionalData=b''):
   """Create a boot option, suppressing duplicates and returning the
   option number.  The filePathList may come from makeHDDevPath,
   makeNICDevPath, makeURIDevPath, unpackBootOption, etc.

   Parameters:
   description: text description
   filePathList: UEFI filePathList
   optionAttr: option attributes
   optionalData: bytes or str (str will be encoded as utf-16-le)
   """
   if isinstance(optionalData, str):
      optionalData = packStr(optionalData)
   opt = {'varAttr':      VAR_ATTRIBS,
          'optionAttr':   optionAttr,
          'description':  description,
          'filePathList': filePathList,
          'optionalData': optionalData}
   return createRawBootOption(packBootOption(opt))

def setOptionalData(optionNumber, value):
   """Change optional data for an existing boot option.  Optional data
   can be either a bytes object (used as-as) or a str (automatically
   converted to UTF-16).  Converting to UTF-16 is potentially useful
   for an option that boots ESXi, but not for all kinds of options.
   """
   option = unpackBootOption(getRawBootOption(optionNumber))
   if isinstance(value, str):
      option['optionalData'] = packStr(value)
   else:
      option['optionalData'] = value
   setRawBootOption(optionNumber, packBootOption(option))

def setDescription(optionNumber, description):
   """Change the description of an existing boot option.
   """
   option = unpackBootOption(getRawBootOption(optionNumber))
   option['description'] = description
   setRawBootOption(optionNumber, packBootOption(option))

def setBootOptionActive(optionNumber, active=True):
   """Set optionNumber to be active (True) or inactive (False).
   """
   option = unpackBootOption(getRawBootOption(optionNumber))
   if active:
      option['optionAttr'] = option['optionAttr'] | LOAD_OPTION_ACTIVE
   else:
      option['optionAttr'] = option['optionAttr'] & ~LOAD_OPTION_ACTIVE
   setRawBootOption(optionNumber, packBootOption(option))

def setBootOptionHidden(optionNumber, hidden=True):
   """Set optionNumber to be hidden (True) or visible (False).
   """
   option = unpackBootOption(getRawBootOption(optionNumber))
   if hidden:
      option['optionAttr'] = option['optionAttr'] | LOAD_OPTION_HIDDEN
   else:
      option['optionAttr'] = option['optionAttr'] & ~LOAD_OPTION_HIDDEN
   setRawBootOption(optionNumber, packBootOption(option))

#
# Creating a boot option for a specific device
#

def getPartUniqueGUID(devfsPath, partNum):
   """Parse GPT header of a disk and return the unique GUID of the
      specified partition.
   """
   GPT_SIG = b'EFI PART'
   with open(devfsPath, 'rb') as rObj:
      rObj.read(512) # Discard LBA0
      lba1 = rObj.read(512)
      if lba1[:8] != GPT_SIG:
         raise ValueError('Disk %s does not contain a GPT partition table'
                          % devfsPath)
      sizeOfEntry = struct.unpack('<I', lba1[84:88])[0]
      # Seek to the partition entry
      rObj.seek(sizeOfEntry * (partNum - 1), 1)
      partLba = rObj.read(sizeOfEntry)
      uniqueGuid = uuid.UUID(bytes_le=partLba[16:32])
      if uniqueGuid.int == 0:
         raise ValueError('Cannot find UUID for partition %d on disk %s'
                          % (partNum, devfsPath))
      return uniqueGuid.hex.upper()

def getOutput(cmd):
   """Execute a command and return stdout as a string.
   """
   rc, out, _ = runcommand.runcommand(cmd, redirectErr=False)
   if rc != 0:
      raise Exception('Command %s failed with status %d' % (cmd, rc))
   return out.decode()

def getHDPartInfoLegacy(diskName, partNum):
   """Get information about a hard drive partition on an ESXi release
      that does not have partedUtil partinfo command, i.e. before 6.7.0.
      Most of the info can be found in partedUtil getptbl, but unique
      partition GUID in the case of GPT is read directly from the disk.
      Only required fields in makeHDDevPath() are populated.
   """
   cmd = ('/bin/partedUtil', 'getptbl', diskName)
   lines = [line for line in getOutput(cmd).split('\n') if line]
   isGpt = (lines[0] == 'gpt')
   for line in lines[2:]:
      # Examples
      # GPT: 1 64 8191 C12A7328F81F11D2BA4B00A0C93EC93B systemPartition 128
      # MBR: 4 32 8191 4 128
      parts = line.split(' ')
      if int(parts[0]) == partNum:
         pi = dict()
         pi['Partition Number'] = parts[0]
         pi['Start sector'] = parts[1]
         pi['End sector'] = parts[2]
         if isGpt:
            pi['Partition Type GUID'] = parts[3]
            pi['Partition Filesystem Type'] = parts[4]
            pi['Partition attributes'] = parts[5]
            pi['Partition Unique GUID'] = getPartUniqueGUID(diskName, partNum)
         else:
            pi['Partition Type'] = parts[3]
            pi['Partition attributes'] = parts[4]
         return pi
   raise ValueError('Partition %d does not exist on disk %s'
                    % (partNum, diskName))

def getEsxVersion():
   """Get ESXi release version in a number tuple.
   """
   verStr = vsi.get('/system/version')['productVersion']
   return [int(n) for n in verStr.split('.')]

def getHDPartInfo(diskName, partNum):
   """Get information from vmkernel about a hard drive partition.
   """
   if uefivar.isVMkernel:
      if getEsxVersion() >= [6, 7, 0]:
         # partinfo is only available on 6.7 and later
         cmd = ('/bin/partedUtil', 'partinfo', diskName, '%d' % partNum)
         out = getOutput(cmd)
         return dict(re.findall(r'([^:]+):\s*(.*)\n', out))
      else:
         return getHDPartInfoLegacy(diskName, partNum)
   else:
      raise NotImplementedError

def makeHDDevPath(diskName, partNum, pathName):
   """Make a short-form device path for hard disk boot.

   Parameters:
   diskName: ESXi pathname of disk to boot from
   partNum: partition number on disk to boot from
   pathName: pathname of bootloader within partition filesystem
   """
   pi = getHDPartInfo(diskName, partNum)
   startLB = int(pi['Start sector'])
   lastLB = int(pi['End sector'])
   sizeLB = lastLB - startLB + 1
   gpt = 'Partition Unique GUID' in pi
   if gpt:
      sig = uuid.UUID(pi['Partition Unique GUID']).bytes_le
   else:
      with open(diskName, 'rb') as f:
         f.seek(MBR_SIGNATURE_OFFSET)
         sig = f.read(MBR_SIGNATURE_LEN)
   dp = packHDDevPath(partNum, gpt, sig, startLB, sizeLB)
   fp = packFPDevPath(pathName)
   return dp + fp + packEndPath()

def getNICInfo(vmnic):
   """Get PCI address and mac address of a NIC from vmkernel.
   """
   props = vsi.get('/net/pNics/%s/properties' % vmnic)
   sbdf = ('{pciSegment:04x}:{pciBus:02x}:{pciSlot:02x}.{pciFunc:x}'.
           format(**props))
   mac = props['macAddr']
   return (sbdf, mac)

def makeNICDevPath(vmnic, macType=1, ipv=None):
   """Make a device path for network boot.  UEFI does not define a
   short-form network boot option, so the PCI path of the NIC is
   always included.

   XXX Experimental.  Unfortunately, creating network boot options is
   a gray area in the UEFI 2.7 spec.  (1) It is unclear whether the
   macType must be 0 or 1, or if either is equally valid.  Testing has
   shown some firmware that accepts macType=0 but ignores and deletes
   the boot option if macType=1 (an AMD Zen Myrtle prototype), and
   some that does exactly the opposite (a VMware VM).  (2) It is
   unclear whether an IPv4 or IPv6 node is required/allowed; this may
   depend on spec version.  (3) Even if an IPv4 or IPv6 node is
   present, firmware may ignore the node and select IPv4 or IPv6 based
   on some other criterion.

   Parameters:
   vmnic: name of nic (e.g., vmnic0)
   macType: type tag to use for mac address (0 or 1)
   ipv: IP version to request, (4, 6, or None)
   """
   (sbdf, mac) = getNICInfo(vmnic)
   fpl = packPCIDevPath(sbdf) + packMACDevPath(mac, macType)
   if ipv == 4:
      fpl += packIPv4DevPath()
   elif ipv == 6:
      fpl += packIPv6DevPath()
   return fpl + packEndPath()

def makeURIDevPath(uri, ipv):
   """Make a short-form device path for URI boot.

   Parameters:
   uri: URI string; must be restricted to ASCII range and %-encoded if needed
   ipv: IP protocol version; must be 4 or 6

   """
   if ipv == 4:
      fpl = packIPv4DevPath()
   elif ipv == 6:
      fpl = packIPv6DevPath()
   else:
      raise ValueError('IP version must be 4 or 6')
   return fpl + packURIDevPath(uri) + packEndPath()

#
# Boot option/order formatting
#

def formatBootOption(optionNumber, verbose=False):
   """Format the specified boot option for printing.
   """
   try:
      raw = getRawBootOption(optionNumber)
   except uefivar.UefiVarNotFound:
      return 'Boot%04X: <<<NONEXISTENT>>>' % optionNumber
   option = unpackBootOption(raw)
   attr = 'inactive'
   oa = option['optionAttr']
   if (oa & LOAD_OPTION_ACTIVE) != 0:
      attr = 'active'
   if (oa & LOAD_OPTION_FORCE_RECONNECT) != 0:
      attr = attr + ', reconnect'
   if (oa & LOAD_OPTION_HIDDEN) != 0:
      attr = attr + ', hidden'
   if (oa & LOAD_OPTION_CATEGORY) == LOAD_OPTION_CATEGORY_APP:
      attr = attr + ', app'
   if verbose:
      out = (bootOptionName(optionNumber) + '\n' +
             pprint.pformat(option) + '\n' +
             'FilePathList: ' + formatDevPath(option['filePathList']) + '\n' +
             'Option attributes: ' + attr + '\n')

      # Try to decode optionalData as UTF-16 and print it.  That is
      #  correct for an option that boots ESXi, but not for all kinds of
      #  options.
      try:
         od = unpackStr(option['optionalData'])
         out = out + 'Optional data interpreted as UTF-16: ' + od + '\n'
      except UnicodeDecodeError:
         pass
      if len(option['optionalData']) == 16:
         out = (out + 'Optional data interpreted as UUID: ' +
                str(uuid.UUID(bytes_le=option['optionalData'])) + '\n')
      return out
   else:
      # Omit '[active]', as that's the normal case.
      return 'Boot%04X: %s%s' % (optionNumber,
                                 option['description'],
                                 (' [%s]' % attr) if attr != 'active' else '')
