#!/usr/bin/env python

import zipfile, struct

class EvilZipStreamWrapper (object):
   def __init__ (self, victim):
      self.victim_fd = victim
      self.position = 0
      self.tells = []
      self.in_file_data = 0

   def tell (self):
      self.tells.append (self.position)
      return self.position

   def seek (self, offset, whence = 0):
      if offset != 0:
         if offset == self.tells[0] + 14:
            # the zipfile module tries to fix up the file header.
            # write Data descriptor header instead,
            # the next write from zipfile
            # is CRC, compressed_size and file_size (as required)
            self.write ("PK\007\010")
         elif offset == self.tells[1]:
            # the zipfile module goes to the end of the file. The next
            # data written definitely is infrastructure (in_file_data = 0)
            self.tells = []
            self.in_file_data = 0
         else:
            raise "unexpected seek for EvilZipStreamWrapper"

   def write (self, data):
      # only test for headers if we know that we're not writing
      # (potentially compressed) data.
      if self.in_file_data == 0:
         if data[:4] == zipfile.stringFileHeader:
            # fix the file header for extra Data descriptor
            hdr = list (struct.unpack (zipfile.structFileHeader, data[:30]))
            hdr[3] |= (1 << 3)
            data = struct.pack (zipfile.structFileHeader, *hdr) + data[30:]
            self.in_file_data = 1
         elif data[:4] == zipfile.stringCentralDir:
            # fix the directory entry to match file header.
            hdr = list (struct.unpack (zipfile.structCentralDir, data[:46]))
            hdr[5] |= (1 << 3)
            data = struct.pack (zipfile.structCentralDir, *hdr) + data[46:]

      self.position += len (data)
      self.victim_fd.write (data)

   def __getattr__ (self, name):
      return getattr (self.victim_fd, name)


if __name__=='__main__':
   import sys, os, traceback

   outfile = sys.stdout

   if len (sys.argv) != 2:
      print >>sys.stderr, "create an ZIP-archive of <dir> to stdout"
      print >>sys.stderr, "   usage: %s <dir> | cat > outfile.zip" % sys.argv[0]
      sys.exit (1)

   src_dir = os.path.abspath (sys.argv[1])

   if not os.path.isdir (src_dir):
      print >>sys.stderr, "create an ZIP-archive of <dir> to stdout"
      print >>sys.stderr, "   usage: %s <dir> | cat > outfile.zip" % sys.argv[0]
      sys.exit (1)

   print >>sys.stderr, """
------------------------------------------------
Trying to use the zipfile module with sys.stdout
------------------------------------------------
"""
   try:
      zfile = zipfile.ZipFile (sys.stdout, 'w', zipfile.ZIP_DEFLATED)
      stripoff = os.path.dirname (src_dir) + os.sep

      for root, dirs, files in os.walk (src_dir):
         for file in files:
            filename = os.path.join (root, file)
            if filename[:len (stripoff)] != stripoff:
               raise RuntimeException, "invalid filename assumptions, please report!"
            zfile.write (filename, filename[len (stripoff):])
      zfile.close ()
      print >>sys.stderr, """
------------------------------------------------------
Worked OK, apparently stdout was redirected to a file.
------------------------------------------------------
"""
   except Exception, e:
      traceback.print_exc ()
      print >>sys.stderr, """
------------------------------------------------
Failed. Now trying with the EvilZipStreamWrapper
------------------------------------------------
"""

      ezfile = EvilZipStreamWrapper (sys.stdout)
      zfile = zipfile.ZipFile (ezfile, 'w', zipfile.ZIP_DEFLATED)
      stripoff = os.path.dirname (src_dir) + os.sep

      for root, dirs, files in os.walk (src_dir):
         for file in files:
            filename = os.path.join (root, file)
            if filename[:len (stripoff)] != stripoff:
               raise RuntimeException, "invalid filename assumptions, please report!"
            zfile.write (filename, filename[len (stripoff):])
      zfile.close ()
      print >>sys.stderr, """
----------
Worked OK.
----------
"""