########################################################################
# Copyright (c) 2016-2023 VMware, Inc.  All rights reserved.           #
########################################################################

import logging
import os
import shutil
import tarfile
import tempfile
import time
from functools import lru_cache

log = logging.getLogger('Ramdisk')

from vmware.runcommand import runcommand, RunCommandError

from . import HostInfo
from .. import Errors
from .Misc import byteToStr, extractTar
from ..Version import Version

RAMDISK_ADD_CMD = '/bin/localcli system visorfs ramdisk add -m %u -M %s ' \
                  '-n %s -p 755 -t %s'
RAMDISK_RM_CMD = '/bin/localcli system visorfs ramdisk remove -t %s'
SECURE_MOUNT_SCRIPT = '/usr/lib/vmware/secureboot/bin/secureMount.py'

def RemoveRamdisk(ramdiskName, target, raiseException=False):
   '''Unmount and remove a ramdisk.
   '''
   def handleError(msg, cause=None):
      if raiseException:
         raise Errors.InstallationError(cause, None, msg)
      else:
         log.warning(msg)

   if os.path.exists(target):
      from esxutils import EsxcliError, runCli
      try:
         ramdisks = runCli(['system', 'visorfs', 'ramdisk', 'list'], True)
      except EsxcliError as e:
         msg = 'Failed to query ramdisk stats: %s' % str(e)
         raise Errors.InstallationError(e, None, msg)
      foundRamdisk = False
      for ramdisk in ramdisks:
         try:
            name, mount = ramdisk['Ramdisk Name'], ramdisk['Mount Point']
         except KeyError as e:
            msg = ('Failed to query ramdisk stats: field "%s" '
                   'not found in output' % e)
            raise Errors.InstallationError(e, None, msg)
         if name == ramdiskName and os.path.samefile(mount, target):
            foundRamdisk = True
            break
         if name == ramdiskName:
            msg = 'Ramdisk %s is not mounted on the target %s' % \
                  (ramdiskName, target)
            handleError(msg)
            return
         if os.path.samefile(mount, target):
            msg = 'Mount point %s mounts a different ramdisk %s' % \
                  (target, name)
            handleError(msg)
            return
      try:
         if not foundRamdisk:
            shutil.rmtree(target)
            return

         for i in range(2, -1, -1):
            # Total 3 tries to remove the ramdisk.
            cmd = RAMDISK_RM_CMD % target
            rc, out = runcommand(cmd)
            if rc == 0:
               break
            log.warning('Failed to remove ramdisk %s: %s, %d retry left',
                        ramdiskName, byteToStr(out), i)
            if i > 0:
               time.sleep(2)
         else:
            handleError('Failed to remove ramdisk %s: %s'
                        % (ramdiskName, byteToStr(out)))
            return
         shutil.rmtree(target)
      except RunCommandError as e:
         handleError('Failed to run %s: %s' % (cmd, e), cause=e)
      except EnvironmentError as e:
         handleError('Cannot remove %s directory: %s' % (target, e),
                     cause=e)

def CreateRamdisk(size, name, target, reserveSize=0):
   """Create and mount a ramdisk.
   """
   if reserveSize < 0 or reserveSize > size:
      raise ValueError('Reserve size should be between 0 and size')

   try:
      os.makedirs(target)
      cmd = RAMDISK_ADD_CMD % (reserveSize, size, name, target)
      rc, out = runcommand(cmd)
   except EnvironmentError as e:
      RemoveRamdisk(name, target)
      msg = 'Failed to create ramdisk %s: %s' % (name, e)
      raise Errors.InstallationError(e, None, msg)
   except RunCommandError as e:
      RemoveRamdisk(name, target)
      msg = ('Failed to run %s: %s' % (cmd, e))
      raise Errors.InstallationError(e, None, msg)

   if rc != 0:
      out = byteToStr(out)
      RemoveRamdisk(name, target)

      cause = None
      msg = 'Failed to create ramdisk %s: %s' % (name, out)
      if 'no space left on device' in out.lower():
         # Embed RamdiskMemoryError, there is not enough memory
         # to reserve.
         cause = Errors.MemoryReserveError(reserveSize,
            'Cannot reserve %u MB of memory for ramdisk %s'
            % (reserveSize, name))
      raise Errors.InstallationError(cause, None, msg)

def GetRamdiskSizeInMiB(ramdiskName):
   """Get size of a ramdisk in MiB.
   """
   from esxutils import EsxcliError, runCli
   try:
      ramdisks = runCli(['system', 'visorfs', 'ramdisk', 'list'], True)
   except EsxcliError as e:
      msg = 'Failed to query ramdisk stats: %s' % str(e)
      raise Errors.InstallationError(e, None, msg)

   try:
      for ramdisk in ramdisks:
         if ramdisk['Ramdisk Name'] == ramdiskName:
            # All sizes output in KiB.
            return ramdisk['Maximum'] / 1024
   except KeyError as e:
      msg = ('Failed to query ramdisk stats: field "%s" not found in output'
             % e)
      raise Errors.InstallationError(e, None, msg)

   raise Errors.InstallationError(None, None,
            'Failed to find ramdisk with name %s' % ramdiskName)

@lru_cache()
def isRamdiskMountSupported():
   """Checks the SECURE_MOUNT_SCRIPT option for ramdisk support."""
   _, out = runcommand(SECURE_MOUNT_SCRIPT)
   # evalutes to true for the new and old interface
   return 'ramdisk' in out.decode()

@lru_cache()
def isIgnoreSigErrOptionAvailable():
   """Check if the --ignoreSigError option is available.

   If this option is not present, the SECURE_MOUNT_SCRIPT doesn't value
   signature violation. The option is required to proceed the mount
   if a violation occurs (e.g., force live mount of VIBs).
   """
   cmd = "%s ramdisk -h" % SECURE_MOUNT_SCRIPT
   _, out = runcommand(cmd)
   return "--ignoreSigError" in out.decode()

def MountTardiskInRamdisk(vibArg, payloadName, tardiskPath, ramdiskName,
                          ramdiskPath, bootPath=None, checkAcceptance=True):
   """Mount and attach a tardisk to an existing ramdisk.
      Parameters:
         vibArg      - VIB ID or the path to the VIB file; secureMount requires
                       this to verify the tardisk
         payloadName - the name of the payload associated with the tardisk
         tardiskPath - local path of a tardisk
         ramdiskName - name of the ramdisk to attach the tardisk
         ramdiskPath - path to the ramdisk
         bootPath    - path to a boot directory (containing imgdb and boot.cfg)
         checkAcceptance - don't mount a tardisk if the signature validation
                           fails
   """

   if isRamdiskMountSupported():
      try:
         log.info('Mount tardisk %s in ramdisk %s', tardiskPath, ramdiskPath)
         curVer = HostInfo.GetEsxVersion()
         if Version.fromstring(curVer) >= Version.fromstring('6.8.8'):
            # secureMount.py API was changed in 6.8.8 with support of operation
            # mode input.
            cmd = '%s ramdisk -v %s -p %s -t %s -r %s' % (
               SECURE_MOUNT_SCRIPT, vibArg, payloadName, tardiskPath,
               ramdiskName)
            if bootPath:
               cmd += ' -b ' + bootPath
            if not checkAcceptance and isIgnoreSigErrOptionAvailable():
               cmd += ' --ignoreSigError'

            rc, out = runcommand(cmd)
         else:
            assert(not bootPath), \
               'Boot path argument is not supported in this release'
            rc, out = runcommand('%s ramdiskMount %s %s %s %s' % (
                                 SECURE_MOUNT_SCRIPT, vibArg, payloadName,
                                 tardiskPath, ramdiskName))
         if rc != 0:
            raise Errors.InstallationError(None, None,
                                           'secureMount returns status %d, '
                                           'output: %s' % (rc, byteToStr(out)))
      except RunCommandError as e:
         log.warning('Failed to execute secureMount: %s', str(e))

   elif not checkAcceptance:
      # Fallback to extraction on failure, this happens during an upgrade
      # where an old secureMount script is present and does not support
      # the ramdiskMount operation.
      # This path doesn't include any secure boot checks and it is the
      # responsibility of the caller to ensure sufficent checks.
      log.info('Fallback to extract tardisk %s', tardiskPath)
      try:
         with tempfile.NamedTemporaryFile() as tmpFd:
            rc, out = runcommand('/bin/vmtar -x %s -o %s'
                                 % (tardiskPath, tmpFd.name))
            log.debug('vmtar returns %d, output: %s', rc, byteToStr(out))
            extractTar(tmpFd.name, ramdiskPath, exception=tarfile.TarError)
      except (RunCommandError, tarfile.TarError) as e:
         msg = 'Failed to extract tardisk %s in ramdisk %s: %s' \
               % (tardiskPath, ramdiskPath, str(e))
         raise Errors.InstallationError(e, None, msg)
   else:
      msg = 'Current ESXi version does not provide a mechanism to mount a ' \
            'tardisk into a ramdisk.'
      raise Errors.InstallationError(None, None, msg)


def UnmountManualTardisk(tardiskName, raiseException=True):
   """Unmount tardisk mounted in tardisks.noauto.
      Such tardisks are mounted to be attached to a ramdisk.
      Parameter:
         tardiskName - filename of the tardisk to be unmounted
   """
   TARDISKS_NOAUTO_PATH = '/tardisks.noauto'

   tardiskPath = os.path.join(TARDISKS_NOAUTO_PATH, tardiskName)
   if os.path.exists(tardiskPath):
      try:
         log.info('Unmounting manual tardisk %s', tardiskPath)
         os.remove(tardiskPath)
      except Exception as e:
         msg = 'Failed to unmount manual tardisk %s: %s' % (tardiskPath, str(e))
         if raiseException:
            raise Errors.InstallationError(e, None, msg)
         else:
            log.warning(msg)
