# Copyright 2019-2021 VMware, Inc.
# All rights reserved. -- VMware Confidential

"""Utilities to setup the ESXi boot partition.

The global imports and UEFI boot option manipulation functions in this module
must maintain compatibility with two previous OnPrem ESXi releases.

Any imports that do not exist on older ESXi releases must be isolated within the
respective functions.
"""

import logging
import os
from tempfile import NamedTemporaryFile

from systemStorage.esxdisk import EsxDisk, iterDisks
from systemStorage.syslinux import installBootSector, patchBootPartition
from systemStorage.vfat import mcopy, mmd

# Relative location of bootloader and syslinux files on ESX.
BOOTLOADER_DIR = os.path.join(os.path.sep, 'usr', 'lib', 'vmware',
                              'bootloader-installer')

log = logging.getLogger(os.path.basename(__file__))
log.setLevel(logging.DEBUG)

def getActiveBootDisk():
   """Return the disk ESX was booted from (assuming it is booted from disk).
   """
   for disk in iterDisks():
      if disk.isActiveBootDisk:
         return disk
   return None

def getUefiArch(path):
   """Return the target architecture of a UEFI (PE) executable.
   """
   MAGIC_OFFSET = 0x3c
   PE_MAGIC = b'PE\x00\x00'
   UEFI_ARCH = {0x8664: 'x64', 0xaa64: 'AA64'}

   with open(path, 'rb') as pe:
      # Read the MAGIC offset
      pe.seek(MAGIC_OFFSET, 0)
      offset = pe.read(1)
      offset = int.from_bytes(offset, 'little')

      # Read the 4-byte PE MAGIC
      pe.seek(offset, 0)
      magic = pe.read(4)
      if magic != PE_MAGIC:
         raise ValueError("%s: not a UEFI executable (bad PE magic)" % path)

      # Read the target architecture (located immediately after the PE MAGIC)
      machineType = pe.read(2)
      machineType = int.from_bytes(machineType, 'little')
      try:
         return UEFI_ARCH[machineType]
      except KeyError:
         raise ValueError("%s: unsuported UEFI machine type %s" %
                          (path, machineType))

def installUefiBootloader(disk, srcRoot, mcopyBin=None, mmdBin=None):
   """Install the UEFI bootloader in the boot partition of the disk, by default
   the boot disk.

   The bootloader payload in esx-base VIB supplies the modules. By default, the
   modules are copied to the boot volume on the host (must be vFAT). In case of
   an upgrade, a prefix to the bootloader tardisk mount path shall be provided
   to find the new modules.
   """
   bootPart = disk.getBootPartition()
   bootPartOffset = bootPart.start * disk.sectorSize

   uefiFiles = [('mboot64.efi', '/EFI/VMware/mboot64.efi'),
                ('crypto64.efi', '/EFI/VMware/crypto64.efi'),
                ('safeboot64.efi', '/EFI/VMware/safeboot64.efi'),
                ('safeboot64.efi', '/EFI/BOOT/BOOT.EFI')]

   for dirPath in ['/EFI', '/EFI/BOOT', '/EFI/VMware']:
      # mmd returns non-zero for folders that exist, we will ignore errors
      # here. Device problem should happen anyway later with mcopy.
      try:
         mmd(disk.path, dirPath, byteOffset=bootPartOffset, exe=mmdBin)
      except RuntimeError:
         log.info('mmd failed to create folder %s on device %s, '
                  'it may already exist', dirPath, disk.path)

   for name, dest in uefiFiles:
      src = os.path.join(srcRoot, name)

      if dest == '/EFI/BOOT/BOOT.EFI':
         dest = '/EFI/BOOT/BOOT%s.EFI' % getUefiArch(src)

      if name == 'crypto64.efi' and not os.path.exists(src):
         # crypto64.efi is packaged starting 7.0 U2.
         log.warn('Skip installing crypto module, source file does not exist')
         continue
      mcopy(disk.path, [src], dest, byteOffset=bootPartOffset, exe=mcopyBin)

def installBiosBootloader(disk, srcRoot, mcopyBin=None):
   """Install BIOS bootloader on a disk, by default operate on the boot disk.

   This installs syslinux on the disk and patches it according to the actual
   ldlinux.sys location. @disk is an EsxDisk object; @srcRoot is the path to
   where the bootloader files are staged.
   """
   # BIOS boot loader is not applicable to disks with sectors > 512
   if disk.sectorSize > 512:
      log.info("Disk %s with sector size %u does not require BIOS boot loader"
               % (disk.name, disk.sectorSize))
      return

   bootPart = disk.getBootPartition()
   bootPartOffset = bootPart.start * disk.sectorSize

   cfg = ['default safeboot.c32 -S 1',
          'nohalt 1']

   with NamedTemporaryFile(mode='w+') as syslinuxCfg:
      syslinuxCfg.write("\n".join(cfg) + "\n")
      syslinuxCfg.flush()
      mcopy(disk.path, [syslinuxCfg.name], '/syslinux.cfg',
            byteOffset=bootPartOffset, exe=mcopyBin)

   for basename in ('mboot.c32', 'safeboot.c32'):
      src = os.path.join(srcRoot, basename)
      mcopy(disk.path, [src], os.path.join(os.sep, basename),
            byteOffset=bootPartOffset, exe=mcopyBin)

   bootCode = os.path.join(srcRoot, 'gptmbr.bin' if disk.isGpt else 'mbr.bin')
   fatBootSectorPath = os.path.join(srcRoot, 'fatBootSector')
   ldlinuxPath = os.path.join(srcRoot, 'ldlinux.sys')

   installBootSector(disk.path, bootCode)
   patchBootPartition(disk.path, bootPart.start, disk.sectorSize,
                      fatBootSectorPath, ldlinuxPath, mcopyBin=mcopyBin)

def addUefiBootDisk(disk=None, setDefault=True):
   """Register a disk as a UEFI boot option.

   On UEFI, create a 'VMware ESXi' boot option for the current installation or
   the installation on a specific disk.

   Any duplicate entry to the same boot partition will be removed. Manual boot
   option creation is required for multiple ESXi across different disks.

   Parameters:
      disk        - disk to create an option for. If None is given, the current
                    boot disk will be used.
      setDefault  - whether to make the created option the first in boot order
                    (default to True).
   """
   from uefi.bootorder import (createBootOption, getBootOrder, getRawBootOption,
                               makeHDDevPath, promoteBootOption, removeBootOption,
                               setBootOrder, unpackBootOption)
   from uefi.uefivar import isUefi

   GPT_BOOTPART_NUM = 1
   BOOTOPTION_DESC = 'VMware ESXi'

   if not isUefi:
      log.info("System is not booted in UEFI mode, "
               "skipping UEFI boot option creation.")
      return

   if disk is None:
      disk = getActiveBootDisk()
      if disk is None:
         log.info("No boot disk found (diskless or PXE), "
                  "skipping UEFI boot option creation.")
         return
      disk.scanPartitions()

   # Partitions must have been scanned.
   assert disk._pt._type is not None

   # A UEFI bug prevents MBR short-form boot option from working properly. And
   # generating a long-form option is impractical. See PR2071443 and PR2101549.
   if not disk.isGpt:
      log.info("Legacy MBR boot disk detected, "
               "skipping UEFI boot option creation.")
      return

   if disk.hasLegacyBootPart:
      # Legacy boot partition was FAT12, which limits filenames to 8.3 format.
      safebootPath = r'\EFI\VMware\safebt64.efi'
   else:
      safebootPath = r'\EFI\VMware\safeboot64.efi'

   try:
      filePathList = makeHDDevPath(disk.path, GPT_BOOTPART_NUM, safebootPath)
      optNum = createBootOption(BOOTOPTION_DESC, filePathList)
      log.info("Created UEFI boot option: [%u] -> %s.", optNum, safebootPath)
   except Exception as e:
      # Boot option creation is not a required step, we could run into firmware issue
      # or out-of-resource problem.
      log.warn("Failed to create UEFI boot option: %s", str(e))
      return

   try:
      if setDefault:
         log.info('Setting UEFI boot option [%u] as default.', optNum)
         promoteBootOption(optNum)
      else:
         # Make sure the option is in the boot order
         bootOrder = getBootOrder()
         if not optNum in bootOrder:
            # Only add the option to the end of boot order if it is newly created,
            # otherwise just leave it where it is.
            bootOrder.append(optNum)
            setBootOrder(bootOrder)

      # Remove legacy and duplicate options.
      dupOpts = []
      for num in getBootOrder():
         opt = unpackBootOption(getRawBootOption(num))
         if (opt['filePathList'] == filePathList or
             opt['description'] == BOOTOPTION_DESC) and num != optNum:
            dupOpts.append(num)

      for num in dupOpts:
         log.info("Removing UEFI boot option [%u] with duplicate name or disk "
                  "path.", num)
         removeBootOption(num, clean=True)
   except Exception as e:
      # Ignore problems of adjusting boot order.
      log.warn("Failed to adjust UEFI boot order: %s", str(e))

def addUefiBootDiskByName(diskName, **kwargs):
   """Set disk pointed to by the given device path as the boot disk.
   """
   disk = EsxDisk(diskName)
   disk.scanPartitions()
   addUefiBootDisk(disk=disk, **kwargs)

def installBootloader(disk=None, srcRoot=None, mcopyBin=None, mmdBin=None):
   """Install the ESX bootloader in the MBR / boot partition.

   This function installs the UEFI bootloader in the boot partition. On x86
   platforms, it also installs the syslinux in the MBR / boot partition.
   """
   if disk is None:
      disk = getActiveBootDisk()
      if disk is None:
         raise OSError("No boot disk found (system is either diskless or "
                       "PXE-booted)")

   disk.scanPartitions()

   if srcRoot is None:
      srcRoot = BOOTLOADER_DIR

   bootModules = os.listdir(srcRoot)
   hasBios = any([name.lower().endswith('.c32') for name in bootModules])
   if hasBios:
      installBiosBootloader(disk, srcRoot, mcopyBin=mcopyBin)

   installUefiBootloader(disk, srcRoot, mcopyBin=mcopyBin, mmdBin=mmdBin)
