#!/usr/bin/python
########################################################################
# Copyright (c) 2010-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.
########################################################################

import json
import logging
import os
import shutil
import sys
import tarfile

import vmkuser
from vmware import runcommand

from .. import Database
from .. import Errors
from .. import Vib
from .. import VibCollection
from . import Installer
from ..Utils import HostInfo
from ..Utils import LockFile
from .InstallerCommon import RemediationResult, StartTransactionResult, \
                             ToolsRamdiskRefreshTrigger

from .. import MIB

LOCKER_ROOT = '/locker/packages/'
PAYLOAD_MAPPING_FILE = os.path.join(LOCKER_ROOT, 'payloadFileMapping.json')

log = logging.getLogger('LockerInstaller')

class Popen(runcommand.VisorPopen):
   '''A Popen class with stdin input and does not redirect stderr.
      Basically a "modern" VisorPopen class tailored to the need.
      We need to emulate the class here for upgrade scenarios where the
      native VisorPopen on host may not have the stdin parameter.
      The class inherits methods from VisorPopen so we can
      use use it with the waitProcessToComplete() method in runcommand.
   '''
   def __init__(self, args, stdin=None):
      '''Class constructor.
         Parameters:
            * args - A command sequence, the first item is the program to
                     execute, and the remaining items are arguments.
            * stdin - stdin file descriptor the subprocess should use.
         Raises:
            * OSError      - If the command cannot be executed.
      '''
      self.returncode = None
      if not stdin:
         stdin = sys.stdin.fileno()

      fromchildfd, toparentfd = os.pipe()
      self.stdout = os.fdopen(fromchildfd, 'rb')

      fromchilderrfd, toparenterrfd = os.pipe()
      self.stderr = os.fdopen(fromchilderrfd, 'rb')

      initfds = [stdin, toparentfd, toparenterrfd]
      try:
         self.pid = vmkuser.ForkExec(args[0], args, None, -1, initfds,
                                     -1, -1, 0)
      except Exception:
         os.close(fromchildfd)
         os.close(fromchilderrfd)
         raise
      finally:
         os.close(toparentfd)
         os.close(toparenterrfd)

class UntarFile(object):
   '''UntarFile class provides write and close methods to untar a tar.gz stream
      to a directory.
   '''
   def __init__(self, dst):
      '''Class constructor.
         Parameters:
            * dst - the directory path to which .tar.gz stream to be extracted.
      '''
      self._dst = dst
      self._cmd = ['/bin/tar', 'xzf', '-', '-C', self._dst]

      tochildfd, fromparentfd = os.pipe()
      try:
         self._stdin = os.fdopen(fromparentfd, 'wb')
         self._p = Popen(self._cmd, stdin=tochildfd)
      except Exception as e:
         os.close(fromparentfd)
         msg = 'Error in creating untar process \'%s\': %s' \
               % (self._cmd, str(e))
         raise Errors.InstallationError(e, None, msg)
      finally:
         os.close(tochildfd)

   def write(self, data):
      self._stdin.write(data)

   def close(self, timeout=120):
      '''Close untar stream and wait for process completion.
         Parameters:
            * timeout - the amount of time in seconds, to wait for output
                        or completion of the process.
      '''
      self._stdin.close()
      try:
         rc, out, err = runcommand.waitProcessToComplete(self._p, self._cmd,
                                                         timeout=timeout)
      except runcommand.RunCommandError as e:
         msg = 'Error while waiting for untar process \'%s\': %s' \
               % (self._cmd, str(e))
         raise Errors.InstallationError(e, None, msg)

      if rc != 0:
         msg = 'Error untarring to %s, return code: %s\n' \
               'stdout: %s\nstderr: %s' % (self._dst, rc, out, err)
         raise Errors.InstallationError(None, None, msg)


LOCKER_POSTINST_TRIGGER_CLASSES = (ToolsRamdiskRefreshTrigger,)

class LockerInstaller(Installer):
   '''LockerInstaller is the Installer class for Locker package files.
      LockerInstaller is only supported on regular booted host. Tools package on
      PXE host is handled by LiveImageInstaller.
      Attributes:
         * database - A Database.Database instance of product locker
         * stagedatabase - Always None, there is no staging support.
   '''
   installertype = 'locker'
   DB_DIR = os.path.join('var', 'db', 'locker')
   DB_LOCKFILE = '/var/run/lockerimgdb.pid'
   priority = 20

   SUPPORTED_VIBS = set([Vib.BaseVib.TYPE_LOCKER, ])
   SUPPORTED_PAYLOADS = set(['tgz', ])
   _SIMULATOR_LOCKER_SIZE_BYTES = 1024 * MIB

   def __init__(self, root=None, **kwargs):
      super(LockerInstaller, self).__init__()

      if root is None:
         self._root = LOCKER_ROOT
      else:
         self._root = root

      # For PXE host, tools payload is handled by LiveImageInstaller only
      if HostInfo.IsPxeBooting():
         raise Errors.InstallerNotAppropriate('Host booted from PXE server or'
               ' there was an error to get boot type. LockerInstaller is not'
               ' supported')
      self.database = Database.Database(os.path.join(self._root, self.DB_DIR),
            dbcreate = True, addprofile=False)
      self.Load()
      self._updateindicator = os.path.join(self._root, 'lockerupdated')
      self.problems = list()
      self._triggers = []

   @property
   def stagedatabase(self):
      return None

   @property
   def root(self):
      return self._root

   def Load(self):
      '''Load locker database.
      '''
      try:
         dbLock = LockFile.acquireLock(self.DB_LOCKFILE)
      except LockFile.LockFileError as e:
         msg = 'Unable to obtain a lock for database I/O: %s' % str(e)
         raise Errors.LockingError(msg)
      try:
         self.database.Load()
      except Exception as e:
         log.warning('Locker DB cannot be loaded: %s' % str(e))
      finally:
         dbLock.Unlock()

   def UpdateVibDatabase(self, newVib):
      """Update missing properties of vib metadata
         Parameters:
            * newvib   - The new vib to use as source
      """
      self.database.vibs[newVib.id].SetSignature(newVib.GetSignature())
      self.database.vibs[newVib.id].SetOrigDescriptor(
                                       newVib.GetOrigDescriptor())

   def SaveDatabase(self):
      """Write out the database of the installer."""
      try:
         dbLock = LockFile.acquireLock(self.DB_LOCKFILE)
      except LockFile.LockFileError as e:
         msg = 'Unable to obtain a lock for database I/O: %s' % str(e)
         raise Errors.LockingError(msg)
      try:
         self.database.Save()
      finally:
         dbLock.Unlock()

   def _GetLockerMetadata(self, imgProfile):
      """Returns locker VIB IDs and associated component objects of the image
         profile.

         Parameters:
            * imgProfile - ImageProfile instance whose vibs and components
                           are collected.

         Returns:
            A tuple of VIB ID set and component collections.
      """
      lockerVibIds = self.GetSupportedVibs(imgProfile.vibs)
      lockerComps = imgProfile.components.GetComponentsFromVibIds(lockerVibIds)
      return lockerVibIds, lockerComps

   def StartTransaction(self, imgprofile, imgstate = None, preparedest = True,
                        forcebootbank = False, **kwargs):
      """Initiates a new installation transaction. Calculate what actions
         need to be taken.

         This method will change product locker

         Parameters:
            * imgprofile  - The ImageProfile instance representing the
                            target set of VIBs for the new image
            * imgstate    - The state of current HostImage, one of IMGSTATE_*
            * preparedest - Boolean, if True, then prepare the destination.
                            Set to false for a "dry run", to avoid changing
                            the destination.
            * forcebootbank - Boolean, if True, skip install of live image
                              even if its eligible for live install
            * stageonly     - If True, do nothing as there is enough space to
                              stage.
         Returns:
            A StartTransactionResult instance that has attributes: installs,
            removes and staged.
            Installs and removes are list of VIB IDs for HostImage.Stage()
            to install to the destination and to remove from the destination,
            in order to make it compliant with imgprofile.

            Staged is always False for LockerInstaller due to absence of
            staging.

            If there is nothing to do, StartTransactionResult(None, None, False)
            is returned.
         Exceptions:
            * InstallationError
      """
      stageonly = kwargs.get('stageonly', False)
      if stageonly:
         msg = 'Stage only is not supported for LockerInstaller.'
         self.problems.append(msg)
         return StartTransactionResult(None, None, False)

      supported, lockerComps = self._GetLockerMetadata(imgprofile)
      curVibIds = set(self.database.vibs.keys())

      keeps = curVibIds & supported
      removes = curVibIds - keeps
      adds = supported - keeps

      if self._platform:
         removes = removes & \
            self.database.vibs.GetPlatformVibIDMap().get(self._platform, set())
         adds = adds & \
            imgprofile.vibs.GetPlatformVibIDMap().get(self._platform, set())

      vibsToAdd = [imgprofile.vibs[vibid] for vibid in adds]
      vibsToRemove = [self.database.vibs[vibid] for vibid in removes]

      if preparedest and (removes or adds):
         # Remove existing VIBs, and commit clean database to be
         # consistent in case the transaction fails.
         self._RemovePayloadFileList(removes)
         self._RemoveVibs(self.database.vibs, removes)
         self.database.Clear()
         self.SaveDatabase()
         self._UnSetUpdated()
         # Add new VIB/component to the database, will not commit until the
         # transaction succeeds.
         for vibid in supported:
            self.database.vibs.AddVib(imgprofile.vibs[vibid])
         self.database.bulletins = lockerComps.GetBulletinCollection()

      # Initiate post-installation triggers
      for cls in LOCKER_POSTINST_TRIGGER_CLASSES:
         self._triggers.append(cls())
      self._UpdateTriggers(vibsToAdd, vibsToRemove)

      return StartTransactionResult(adds, removes, False)

   def OpenPayloadFile(self, vibid, payload, read = False, write = True,
                       isBaseMisc=False):
      '''Creates and returns a File-like object for writing to a given payload.
         Only write is supported.

         Parameters:
            * vibid   - The Vib id containing the payload
            * payload - The Vib.Payload instance to read or write
            * read    - Must be False; ready is not supported
            * write   - Set to True to get a File object for writing
                        to the payload.
            * isBaseMisc - Not used, defaults to False.
         Returns:
            A File-like object, must support write and close methods.
            None if the desired read/write is not supported.
         Exceptions:
            AssertionError    - neither read nor write is True, or both are true
            InstallationError - Cannot open file to write or read
      '''
      if isBaseMisc:
         raise NotImplementedError("Do not know how to handle isBaseMisc")

      Installer.OpenPayloadFile(self, vibid, payload, read, write)
      if write == True:
         if payload.payloadtype not in self.SUPPORTED_PAYLOADS:
            log.debug("Payload %s of type '%s'  in VIB '%s' is not supported by "
                  "LockerInstaller." % (payload.name, payload.payloadtype,
                         vibid))
            return None
         return UntarFile(self._root)
      else:
         raise ValueError('OpenPayloadFile for read is not supported in '
               'LockerInstaller')

   def Cleanup(self, **kwargs):
      '''Clean up locker packages directory. Since there is no space for
         staging, locker packages content will be cleaned.
      '''
      self.database.Clear()
      try:
         shutil.rmtree(self.database.dbpath)
         shutil.rmtree(self._root)
         os.makedirs(self._root)
      except Exception as e:
         log.warning('There was an error in cleaning up product locker: %s'
                     % e)

   def CompleteStage(self):
      '''Complete the staging of live image by writing out the database.
      '''
      self.SaveDatabase()
      self._SetUpdated()

      # Invoke post install triggers here as LockerInstaller only stages.
      self._RunTriggers(ignoreErrors=True)

   def Remediate(self, checkmaintmode=True, **kwargs):
      """Nothing to do here, as there is no space to stage in Locker.

         Returns:
            A Boolean, always False, as a reboot is not needed.
         Exceptions:
            * HostNotChanged - If host is not changed in previous Stage command.
      """
      if os.path.exists(self._updateindicator):
         self._UnSetUpdated()
         return RemediationResult(False, None)
      else:
         raise Errors.HostNotChanged('Locker files not chaged.')

   def UpdateImageProfile(self, newImageProfile):
      """Update component in the database using new image profile.
      """
      newLockerVibs, newLockerComps = self._GetLockerMetadata(newImageProfile)
      if self._platform:
         if (newLockerVibs != self.database.vibs.keys() or
             set(newLockerComps.GetComponentIds()) !=
             self.database.bulletins.keys()):
            self.database.vibs = VibCollection.VibCollection()
            for vId in newLockerVibs:
               self.database.vibs.AddVib(newImageProfile.vibs[vId])
            self.database.bulletins = newLockerComps.GetBulletinCollection()
            self.SaveDatabase()
      elif (set(newLockerComps.GetComponentIds()) !=
          self.database.bulletins.keys()):
         self.database.bulletins = newLockerComps.GetBulletinCollection()
         self.SaveDatabase()

   def _RemovePayloadFileList(self, vibIDs):
      '''Remove the file-list mapping for individual payloads for given list
         of vibIDs
      '''
      # If the payload mapping file doesn't exist, nothing to do here.
      if not os.path.exists(PAYLOAD_MAPPING_FILE):
         return

      payloadFileDict = dict()
      try:
         with open(PAYLOAD_MAPPING_FILE, 'r') as f:
            payloadFileDict = json.load(f)
      except Exception as e:
         msg = 'Unable to read file %s: %s' % (PAYLOAD_MAPPING_FILE, str(e))
         raise Errors.InstallationError(e, vibIDs, msg)

      # Go through the list of vibs getting removed and delele the entries
      # from the dictionary if they exist.
      for vibid in vibIDs:
         vib = self.database.vibs[vibid]
         if vib.name in payloadFileDict:
            del payloadFileDict[vib.name]

      # If dictionary is now empty, we need to remove the file otherwise we
      # will dump the updated dictionary into the file.
      if not payloadFileDict:
         os.remove(PAYLOAD_MAPPING_FILE)
      else:
         try:
            with open(PAYLOAD_MAPPING_FILE, 'w') as fp:
               json.dump(payloadFileDict, fp)
         except Exception as e:
            msg = 'Unable to write to file %s: %s' \
                  % (PAYLOAD_MAPPING_FILE, str(e))
            raise Errors.InstallationError(e, vibIDs, msg)

   def _SavePayloadFileList(self, payload, vib, sourcefp):
      '''Save the file-list for individual payloads so that locker vibs with
         multiple payloads can be recreated based on this list.
      '''
      payloadFileDict = dict()
      # If payload file mapping already exists, then we load the dictionary
      # otherwise we will end up deleting information for other vibs.
      if os.path.exists(PAYLOAD_MAPPING_FILE):
         try:
            with open(PAYLOAD_MAPPING_FILE, 'r') as f:
               payloadFileDict = json.load(f)
         except Exception as e:
            msg = 'Unable to read file %s: %s' % (PAYLOAD_MAPPING_FILE, str(e))
            raise Errors.InstallationError(e, [vib.id], msg)

      if not payload.payloadtype in self.SUPPORTED_PAYLOADS:
         msg = 'Payload %s is not supported in LockerInstaller' \
               % payload.name
         log.error(msg)
         return
      with tarfile.open(mode='r:gz', fileobj=sourcefp,
                        format=tarfile.GNU_FORMAT) as tarfObj:
         payloadFileDict.setdefault(vib.name, dict())[payload.name] = \
            tarfObj.getnames()

      sourcefp.seek(0)

      try:
         with open(PAYLOAD_MAPPING_FILE, 'w') as fp:
            json.dump(payloadFileDict, fp)
      except Exception as e:
         msg = 'Unable to write to file %s: %s' \
               % (PAYLOAD_MAPPING_FILE, str(e))
         raise Errors.InstallationError(e, [vib.id], msg)

   def _CheckDiskSpaceForVib(self, arvib):
      '''Compute needed space of the vib, and check if the locker parition
         has enough space.
         When the function is called, we expect the VIB is for locker and
         all of the payloads are of tgz type, the only supported one.
         Parameter:
            arvib - ArFileVib object with the ar file opened for read
      '''
      LOCKER_PATH = os.path.join(os.path.sep, 'locker')

      if not arvib.vibtype in self.SUPPORTED_VIBS:
         raise ValueError('VIB %s is not supported in LockerInstaller'
                          % arvib.id)

      # Get current available size
      # If simulator, then return a hardcoded value for free space available
      # in locker partition
      if HostInfo.HostOSIsSimulator():
         spaceAvail = self._SIMULATOR_LOCKER_SIZE_BYTES
      else:
         spaceAvail = HostInfo.GetFsFreeSpace(LOCKER_PATH)
      log.debug('Locker partition has %d bytes of free space' % spaceAvail)

      # Calculate unpacked size using file members of all tgz archives
      neededSize = 0
      for payload, sourcefp in arvib.IterPayloads():
         if not payload.payloadtype in self.SUPPORTED_PAYLOADS:
            raise ValueError('Payload %s of VIB %s is not supported in '
                             'LockerInstaller' % (payload.name, arvib.id))
         with tarfile.open(mode='r:gz', fileobj=sourcefp,
                           format=tarfile.GNU_FORMAT) as tarfObj:
            for tarInfo in tarfObj.getmembers():
               if tarInfo.isfile():
                  neededSize += tarInfo.size
         log.debug('Payload %s of VIB %s requires %d bytes after extraction'
                   % (payload.name, arvib.id, neededSize))

      # Previous VIB in locker is already removed in StartTransaction(), simply
      # check free space is enough to fit this VIB with 10MB buffer added.
      neededSize += 10 * MIB
      availMiB = int(round(spaceAvail / MIB, 0))
      neededMiB = int(round(neededSize / MIB, 0))
      if neededMiB > availMiB:
         msg = 'VIB %s requires %d MB free space in the locker partition to ' \
               'be safely installed, however the parition only has %d MB ' \
               'available. Please clean up the locker partition and retry ' \
               'the transaction.' % (arvib.name, neededMiB, availMiB)
         raise Errors.InstallationError(None, [arvib.id], msg)

   def _SetUpdated(self):
      if not os.path.isfile(self._updateindicator):
         f = open(self._updateindicator, 'w')
         f.close()

   def _UnSetUpdated(self):
      if os.path.exists(self._updateindicator):
         os.unlink(self._updateindicator)

   def _RemoveVibs(self, allvibs, removes):
      for vibid in removes:
         log.debug('Removing files from productLocker for VIB %s' % (vibid))
         for filepath in allvibs[vibid].filelist:
            filepath = filepath.lstrip('/')
            realpath = os.path.join(self._root, filepath)
            if os.path.isfile(realpath):
               try:
                  os.unlink(realpath)
               except EnvironmentError as e:
                  log.warning('Unable to delete file [%s]: %s, skipping...' %
                        (realpath, e))

   def _UpdateTriggers(self, adds, removes):
      '''Activate trigger for VIB installation or removal
      '''
      for vib in adds:
         for trigger in self._triggers:
            trigger.Match(vib, 'add')
      for vib in removes:
         for trigger in self._triggers:
            trigger.Match(vib, 'remove')

   def _RunTriggers(self, ignoreErrors=False):
      '''Execute post installation triggers.
      '''
      for trigger in self._triggers:
         try:
            trigger.PostInstallAction()
         except Exception as e:
            if ignoreErrors:
               log.warning("Error running %s: %s", trigger.NAME, str(e))
            else:
               raise

