#!/usr/bin/python

'''This module provides a class for manipulating archives created with GNU ar.
It provides support for the GNU extension for long file names, but it currently
only supports extraction.'''

from __future__ import print_function

import os
from stat import *
import struct
import time

from .Misc import byteToStr, seekable

AR_FILEMAGIC = b'\140\012'
AR_FILEMAGIC_STR = '\140\012'

class ArError(Exception):
   def __init__(self, filename, msg):
      Exception.__init__(self, filename, msg)
      self.filename = filename
      self.msg = msg

def extractHeaderFromInfo(info, length):
   """Generate the header info from ArInfo object.
   """
   header = '%-16.16s%-12.12s%-6.6s%-6.6s%-8.8s%-10.10s%-2.2s' % (
      info.filename, info.timestamp, info.uid, info.gid, oct(info.mode),
      length, AR_FILEMAGIC_STR)
   return header

class ArInfo(object):
   '''A class representing a member of the archive.  It contains the following
      members:
      * filename  - the original name of the member (string)
      * timestamp - the mtime of the member (int)
      * uid       - UID of file (int)
      * gid       - GID of file (int)
      * mode      - permissions mask of file (int)
      * size      - length of file data
      * offset    - offset of archive where file data begins
   '''
   def __init__(self, filename, timestamp=0, uid=0, gid=0, mode=0, size=0,
                offset=0):
      self.filename = filename
      self.timestamp = timestamp
      self.uid = uid
      self.gid = gid
      self.mode = mode
      self.size = size
      self.offset = offset

class _FileInFile(object):
   """A thin wrapper around an existing file object that
      provides a part of its data as an individual file
      object.  Taken from python 2.5's tarfile module.
   """

   BUFFER_SIZE = 2048

   def __init__(self, fileobj, size, offset = None):
      self.fileobj = fileobj
      self.size = size
      self.position = 0
      self.offset = offset
      self._seekable = self.offset is not None and seekable(fileobj)

   def close(self):
      """Skip to the end of the file"""
      if self._seekable:
          self.fileobj.seek(self.offset + self.size)
      else:
          while self.position < self.size:
             self.read(self.BUFFER_SIZE)

   def tell(self):
      """Return the current file position.
      """
      if self._seekable:
         return self.position
      else:
         msg = "Underlying file object does not support seek."
         raise NotImplementedError(msg)

   def seekable(self):
      """Return true if the underlying file descriptor supports .seek().
      """
      return self._seekable

   def seek(self, position):
      """Seek to a position in the file.
      """
      if self._seekable:
         if 0 <= position <= self.size:
            self.position = position
         else:
            raise ValueError("Seek(%d) exceeds bounds [0, %d]" %
                             (position, self.size))
      else:
         msg = "Underlying file object does not support seek."
         raise NotImplementedError(msg)

   def read(self, size=None):
      """Read data from the file.
      """
      if self._seekable:
         self.fileobj.seek(self.offset + self.position)
      if size is None:
         size = self.size - self.position
      else:
         size = min(size, self.size - self.position)
      bytesread = self.fileobj.read(size)
      self.position += len(bytesread)
      return bytesread

class ArFile(object):
   """A class representing a GNU ar archive. Supports the GNU extension to
      handle long file names when reading, but can only write short filenames.
   """

   def __init__(self, name=None, mode='rb', fileobj=None):
      """Opens fileobj for reading as a stream containing an ar archive.
            Parameters:
               * name    - A string specifying the name of an ar file.
               * mode    - One of 'r', 'r+', or 'w' to open a file for reading,
                           updating, or writing.
               * fileobj - A file or file-like object supporting at least a
                           read() method.
            Raises:
               * ArFileError - If format is invalid.
               * IOError     - On an error attempting to read or write.
         Notes:
            One of name or fileobj must be specified. If fileobj is specified,
            it does not need to support a seek() method. However, if it does
            not support seeking, use of the object will be limited to
            iteration. Attempting to call any methods other than Close() will
            cause an exception.
      """
      # Just force binary mode. Things will get fouled up without it.
      if not mode.endswith("b"):
         mode += "b"

      if name:
         self._fileobj = open(name, mode)
         self._filename = name
      elif fileobj:
         if mode.startswith('w'):
            raise ValueError('Must initialize with name in w, wb modes')
         self._fileobj = fileobj
         if hasattr(fileobj, 'name'):
            self._filename = fileobj.name
         else:
            self._filename = str(fileobj)
      else:
         raise ValueError('Must initialize with either name or file object')

      self._seekable = seekable(self._fileobj)

      self._memberfile = None
      self._longfns = b''

      if mode.startswith('w'):
         self._fileobj.write(b'!<arch>\n')
         self.filelist = list()
      elif mode.startswith("r"):
         armagic = self._fileobj.read(8)
         if armagic != b'!<arch>\n':
            if name:
               self._fileobj.close()
            raise ArError(self._filename, 'Not a valid ar file (magic=%s)'
                          % armagic)
         if self._seekable:
            self.filelist = list()
            info = self._parseHeader()
            while info:
               self.filelist.append(info)
               self._fileobj.seek(info.size, 1)
               info = self._parseHeader()
      else:
         if name:
            self._fileobj.close()
         raise ArError(self._filename, "Unsupported mode '%s'." % mode)

   def seekable(self):
      """Return True if the underlying file object supports .seek().
      """
      return self._seekable

   def Close(self):
      '''Close the archive.  Further attempts to read member data will fail.'''
      self._fileobj.close()

   def _parseHeader(self):
      def bufferToStr(buf):
         return byteToStr(buf.strip(b' \00'))

      def intfromstr(x, radix=10):
         x = x.strip(b' \00')
         if x: return int(x, radix)
         else: return 0

      header = self._fileobj.read(1)
      #
      # Some versions of ar insert '\n' at end of last data section so
      # entire header is shifted down by one.  We check file magic
      # to align the header string. (from libbb/unarchive.c)
      #
      if header == b'\n':
         header = self._fileobj.read(60)
      else:
         header += self._fileobj.read(59)

      if not header:
         return None
      elif len(header) != 60:
         raise ArError(self._filename, 'Unexpected EOF')

      (filename, timestamp, uid, gid,
       mode, size, filemagic) = struct.unpack("16s12s6s6s8s10s2s", header)

      if filemagic != AR_FILEMAGIC:
         raise ArError(self._filename, 'Unexpected file magic')

      try:
         timestamp = intfromstr(timestamp)
         uid = intfromstr(uid)
         gid = intfromstr(gid)
         mode = intfromstr(mode, 8)
         size = intfromstr(size)
      except ValueError:
         raise ArError(self._filename, 'Invalid format: timestamp=%s '
                       'uid=%s gid=%s mode=%s size=%s header=[%s]' %
                       (timestamp, uid, gid, mode, size, header))

      filename = bufferToStr(filename)

      if filename == "/":
         # We don't deal with symbol tables; go to the next header.
         if self._seekable:
            self._fileobj.seek(size, 1)
         else:
            self._fileobj.read(size)
         return self._parseHeader()
      if filename == "//":
         # this is a long file name index. Fill in self._longfns, then go to
         # the next header.
         self._longfns = self._fileobj.read(size)
         return self._parseHeader()

      if self._longfns:
         try:
            start = int(filename.lstrip('/'))
            end = self._longfns.find(b'\n', start)
            filename = self._longfns[start:end]
            filename = bufferToStr(filename)
         except (ValueError, IndexError):
            raise ArError(self._filename, 'Error getting long filename')

      filename = filename.rstrip('/')

      if self._seekable:
         offset = self._fileobj.tell()
      else:
         offset = None

      return ArInfo(filename, timestamp, uid, gid, mode, size, offset)

   def next(self):
      """A generator method that iterates through the ar stream
         one header and member at a time.  In order to align itself
         with the next header, this method will read to the end of the
         current member.  This assumes that no one has manipulated
         the _fileobj attribute of this class directly.
         Returns:
            A (ArInfo, memberfile) tuple, where memberfile is
            an instance of _FileInFile, a file object representing
            the ar file member corresponding to ArInfo.  memberfile can
            itself be used to stream the contents from the ar member, and
            supports tell(), read(), and close() methods.
      """
      # Seek to the end of the current member
      if self._memberfile:
         self._memberfile.close()
      elif self._seekable:
         # If this is our first time, rewind to the beginning.
         self._fileobj.seek(8, 0)
      info = self._parseHeader()
      if not info:
         raise StopIteration
      self._memberfile = _FileInFile(self._fileobj, info.size, info.offset)
      return info, self._memberfile

   __next__ = next

   def __iter__(self):
      # If someone asks for a new iterator, start from the beginning.
      if self._seekable:
         self._memberfile = None
         self._fileobj.seek(0, 0)
      return self

   def NameToInfo(self, filename):
      '''Returns ArInfo object corresponding to filename, or None
         if no archive member matches filename.
      '''
      if not self._seekable:
         msg = "NameToInfo method requires file object supporting seek."
         raise ArError(self._filename, msg)

      for i in self.filelist:
         if filename == i.filename:
            return i
      return None

   def Read(self, member):
      '''Read file data of archive member.  Member may be either a file name
         or an instance of ArInfo.  Entire data is returned as a string. '''
      if not self._seekable:
         msg = "Read method requires file object supporting seek."
         raise ArError(self._filename, msg)

      if isinstance(member, ArInfo):
         info = member
      else:
         info = self.NameToInfo(member)
         if not info:
            raise KeyError('Archive member not found: %s' % member)

      self._fileobj.seek(info.offset, 0)
      return self._fileobj.read(info.size)

   def Extractfile(self, member):
      '''Returns a file object for a member of the archive.
         member may be either a file name or an ArInfo instance.
         The file object is read-only and provides these methods:
         read(), seek(), tell()
      '''
      if not self._seekable:
         msg = "Extractfile method requires file object supporting seek."
         raise ArError(self._filename, msg)

      if isinstance(member, ArInfo):
         info = member
      else:
         info = self.NameToInfo(member)
         if not info:
            raise KeyError('Archive member not found: %s' % member)
      return _FileInFile(self._fileobj, info.size, info.offset)

   def Writestr(self, info, bytes):
      '''Write ar header and data bytes at current file pos.
         info can be ArInfo object or archive member name.
         If a name, default values will be filled in. '''
      if not self._seekable:
         msg = "Writestr method requires file object supporting seek."
         raise ArError(self._filename, msg)

      if not isinstance(info, ArInfo):
         info = ArInfo(filename=info, timestamp=int(time.time()),
                       mode=0o644)
      if len(info.filename) > 15:
         raise ArError(self._filename, "Long filename writes are "
                       "not currently supported.")

      header = extractHeaderFromInfo(info, len(bytes))
      self._fileobj.write(header.encode('utf-8'))
      self._fileobj.write(bytes)
      #
      # GNU ar apparently expects short word (2-byte)
      # alignment of headers, and \n is used to pad if needed
      #
      if self._fileobj.tell() & 0x01:
         self._fileobj.write(b'\n')
      self.filelist.append(info)

   def Writefile(self, info, fileobj):
      '''Write ar header and data bytes at current file pos.
         info must be ArInfo object.'''
      if not self._seekable:
         msg = "Writefile method requires file object supporting seek."
         raise ArError(self._filename, msg)

      if not isinstance(info, ArInfo):
         raise ArError(self._filename, "ArInfo must be provided")
      if len(info.filename) > 15:
         raise ArError(self._filename, "Long filename writes are "
                       "not currently supported.")

      header = extractHeaderFromInfo(info, info.size)
      self._fileobj.write(header.encode('utf-8'))
      readbytes = info.size
      while readbytes > 0:
         buf = fileobj.read(min(64 * 1024, readbytes))
         readbytes -= len(buf)
         self._fileobj.write(buf)
      #
      # GNU ar apparently expects short word (2-byte)
      # alignment of headers, and \n is used to pad if needed
      #
      if self._fileobj.tell() & 0x01:
         self._fileobj.write(b'\n')
      self.filelist.append(info)

   def Write(self, filename):
      '''Append filename to this archive'''
      if not self._seekable:
         msg = "Write method requires file object supporting seek."
         raise ArError(self._filename, msg)

      stats = os.stat(filename)
      info = ArInfo(os.path.basename(filename), stats[ST_MTIME],
                    stats[ST_UID], stats[ST_GID], stats[ST_MODE])
      with open(filename, 'rb') as fp:
         self.Writestr(info, fp.read())

if __name__ == '__main__':
   import sys

   def usage():
      sys.stderr.write('Usage: %s [rxt] archive [member]\n'
                       '  r - replace or add member to archive\n'
                       '  x - extract member from archive (or all if omitted)\n'
                       '  t - list members in archive\n'
                       % os.path.basename(sys.argv[0]))
      raise SystemExit()

   if len(sys.argv) < 3:
      usage()
   (op, archive) = sys.argv[1:3]
   if len(sys.argv) < 4:
      member = None
   else:
      member = sys.argv[3]

   if op == 'r':
      if not member:
         usage()
      if os.path.exists(archive):
         a = ArFile(archive, 'r+b')
      else:
         a = ArFile(archive, 'wb')
      a.Write(member)
      a.Close()

   elif op == 'x':
      a = ArFile(archive, 'rb')
      if member:
         f = open(member, 'wb')
         f.write(a.Read(member))
         f.close()
      else:
         for info in a.filelist:
            f = open(info.filename, 'wb')
            f.write(a.Read(info))
            f.close()
      a.Close()

   elif op == 't':
      a = ArFile(archive, 'rb')
      for info in a.filelist:
         print(info.filename)
      a.Close()
