#!/usr/bin/python
########################################################################
# Copyright (C) 2010,2020 VMware, Inc.                                 #
# All Rights Reserved                                                  #
########################################################################

import os
import posixpath
import tarfile

from .Utils import EsxGzip
from .Utils.Misc import isString

class PayloadTar(object):
   """Implements a representation of a gzipped tar which can be used to cache
      non-tardisk payloads for use by the installer.
   """
   PAYLOADTAR_PREFIX = "var/db/payloads"

   def __init__(self, filepath):
      """Class constructor.
            Parameters:
               * filepath - Must be either a string, giving the path of a file
                            to write the gzipp'ed tar output, or a file-like
                            object where gzipp'ed tar output will be written.
                            In either case, the file or file-like object must
                            be writable.
      """
      if isString(filepath):
         self.fobj = EsxGzip.GzipFile(filepath, "wb")
      else:
         self.fobj = EsxGzip.GzipFile(fileobj=filepath, mode="wb")

      self.tarfile = tarfile.open(fileobj=self.fobj, mode="w")

   def _AddFile(self, dirPath, fileName, size, fObj):
      """Add a file into the tar, create the folder is not present in the tar.
      """
      self._AddDirectory(dirPath)
      fullPath = posixpath.join(dirPath, fileName)
      tarinfo = tarfile.TarInfo(fullPath)
      tarinfo.mode = 0o644
      tarinfo.size = size
      self.tarfile.addfile(tarinfo, fObj)

   def AddPayload(self, fobj, payload, payloadfn=None):
      if payloadfn is None:
         payloadfn = payload.name
      dirPath = posixpath.join(self.PAYLOADTAR_PREFIX, payload.payloadtype)
      self._AddFile(dirPath, payloadfn, payload.size, fobj)

   def _AddDirectory(self, directorypath):
      dirparts = directorypath.split("/")
      for i in range(len(dirparts)):
         dirname = "/".join(dirparts[:i+1])
         if not dirname or dirname in self.tarfile.getnames():
            continue
         tarinfo = tarfile.TarInfo(dirname)
         tarinfo.type = tarfile.DIRTYPE
         tarinfo.mode = 0o755
         self.tarfile.addfile(tarinfo)

   def FromDirectory(self, dirPath):
      """Populate this PayloadTar from a directory, the payload files should
         be sorted by their types in the directory.
      """
      for root, _, files in os.walk(dirPath, topdown=True):
         if files:
            plType = os.path.basename(root)
            for f in files:
               tarDir = os.path.join(self.PAYLOADTAR_PREFIX, plType)
               fullPath = os.path.join(root, f)
               size = os.stat(fullPath).st_size
               with open(fullPath, 'rb') as fObj:
                  self._AddFile(tarDir, f, size, fObj)

   def Close(self):
      self.tarfile.close()
      self.fobj.close()

   close = Close
