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

"""Read/write/list/remove UEFI variables.

Works on sufficiently recent versions of either ESXi or Linux.  On
ESXi, uses the VSI bindings for UEFI Runtime Services.  On Linux, uses
the efivars filesystem, which is assumed to be mounted in the standard
place, /sys/firmware/efi/efivars.

API conventions:

Each name is a str consisting of the variable's VariableName (decoded
from UTF-16 into a Python string) and its VendorGuid, separated by a
hyphen.  The GUID is in standard GUID print format, with hex digits in
lower case. Example: BootOrder-8be4df61-93ca-11d2-aa0d-00e098032b8c.

Each value is a bytes object consisting of the variable's UINT32
Attributes, formatted as 4 bytes in little-endian order, followed by
its Data.
"""

import os

# Check system type
isVMkernel = (os.uname().sysname == 'VMkernel')
if isVMkernel:
   from vmware import vsi
   isUefi = vsi.get('/hardware/firmwareType') == 1
else:
   isUefi = os.path.exists('/sys/firmware/efi')

class UefiVarError(Exception):
   """Generic error accessing a UEFI variable
   """
   pass

class UefiVarNotFound(UefiVarError):
   """UEFI variable not found on get or remove
   """
   pass

EfiGlobalVariableGUID = '8be4df61-93ca-11d2-aa0d-00e098032b8c'
EfiImageSecurityDatabaseGUID = 'd719b2cb-3d3a-4596-a3bc-dad00e67656f'

# Attributes
EfiVariableNonVolatile = 0x00000001
EfiVariableBootserviceAccess = 0x00000002
EfiVariableRuntimeAccess = 0x00000004
EfiVariableHardwareErrorRecord = 0x00000008
EfiVariableAuthenticatedWriteAccess = 0x00000010
EfiVariableTimeBasedAuthenticatedWriteAccess = 0x00000020
EfiVariableAppendWrite = 0x00000040
EfiVariableEnhancedAuthenticatedAccess = 0x00000080

esxPath = '/uefi/variable/'
linPath = '/sys/firmware/efi/efivars/'

esxVarInfoPath = '/uefi/variableInfo/attributes/'

def list():
   """List the names of all defined UEFI variables.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')
   if isVMkernel:
      return vsi.list(esxPath)
   else:
      return os.listdir(linPath)

def get(name):
   """Return the value of a UEFI variable, given its name.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')
   try:
      if isVMkernel:
         return vsi.get(esxPath + name)
      else:
         with open(linPath + name, 'rb') as f:
            return f.read()
   except OSError as e:
      # Linux case
      if e.args[0] in (os.errno.EINVAL, os.errno.ENOENT):
         raise UefiVarNotFound from e
      else:
         raise UefiVarError(e.args[1]) from e
   except Exception as e:
      # vmkernel case; ugly because vmware.vsi doesn't subclass Exception
      if type(e) is Exception:
         if e.args[0] == 'Not found':
            raise UefiVarNotFound from e
         else:
            raise UefiVarError(e.args[0]) from e
      else:
         raise

def set(name, value):
   """Set the value of a UEFI variable, given its name.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')
   try:
      if isVMkernel:
         vsi.set(esxPath + name, value)
      else:
         with open(linPath + name, 'wb') as f:
            f.write(value)
   except OSError as e:
      # Linux case
      raise UefiVarError(e.args[1]) from e
   except Exception as e:
      # vmkernel case; ugly because vmware.vsi doesn't subclass Exception
      if type(e) is Exception:
         raise UefiVarError(e.args[0]) from e
      else:
         raise

def remove(name):
   """Remove a UEFI variable, given its name.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')
   try:
      if isVMkernel:
         vsi.set(esxPath + name, b'\x00\x00\x00\x00')
      else:
         os.unlink(linPath + name)
   except OSError as e:
      # Linux case
      if e.args[0] in (os.errno.EINVAL, os.errno.ENOENT):
         raise UefiVarNotFound from e
      else:
         raise UefiVarError(e.args[1]) from e
   except Exception as e:
      # vmkernel case; ugly because vmware.vsi doesn't subclass Exception
      if type(e) is Exception:
         if e.args[0] == 'Not found':
            raise UefiVarNotFound from e
         else:
            raise UefiVarError(e.args[0]) from e
      else:
         raise

def getSpace(attr):
   """Get remaining variable storage space.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')

   if isVMkernel:
      vsiPath = esxVarInfoPath + str(attr)
      return vsi.get(vsiPath)['remainingVarStorageSize']
   else:
      raise UefiVarError('Not implemented.')

def getMaxVarSize(attr):
   """Get maximum variable size.
   """
   if not isUefi:
      raise UefiVarError('System not booted in UEFI mode')

   if isVMkernel:
      vsiPath = esxVarInfoPath + str(attr)
      return vsi.get(vsiPath)['maximumVarSize']
   else:
      raise UefiVarError('Not implemented.')

