import tarfile

from cStringIO import StringIO

from util.registry.tarlayerformat import TarLayerFormat
from util.registry.gzipwrap import GzipWrap

class TarfileAppender(TarLayerFormat):
  """ Helper class which allows for appending entries to a gzipped-tarfile and doing so
      in a streaming manner.
  """
  def __init__(self, base_tar_file, entries):
    super(TarfileAppender, self).__init__(self._get_tar_iterator)
    self.entries = entries
    self.base_tar_file = base_tar_file
    self.first_info = None

  def get_stream(self):
    return GzipWrap(self.get_generator())

  def after_tar_layer(self, current_layer):
    pass

  def check_tar_info(self, tar_info):
    if not self.first_info:
      self.first_info = tar_info
    return True

  def _get_tar_iterator(self):
    # Yield the contents of the base tar.
    yield self.base_tar_file

    # Construct an in-memory tar containing the entries to append, and then yield
    # its data.
    def add_entry(arch, dir_path, contents=None):
      info = tarfile.TarInfo(dir_path)
      info.uid = self.first_info.uid
      info.gid = self.first_info.gid
      info.mode = self.first_info.mode
      info.mtime = self.first_info.mtime

      info.type = tarfile.REGTYPE if contents else tarfile.DIRTYPE

      if contents:
        info.size = len(contents)

      arch.addfile(info, fileobj=StringIO(contents) if contents else None)

    append_tarball = StringIO()
    with tarfile.open(fileobj=append_tarball, mode='w') as updated_archive:
      for entry in self.entries:
        add_entry(updated_archive, entry, self.entries[entry])

    # To make tarfile happy.
    append_tarball.seek(0)
    yield append_tarball