#!/usr/bin/python3

import os
import argparse
import zipfile
import zlib # For calculating crc32
import types # For binding seek and tell to ZipExtFile

def seekZipExtFile(self, offset, from_what = 0):
    """Seek method for ZipExtFile
    
    Allows repositioning within a zipfile. Does not perform actual reposition,
    but rather reads forward if the seek is ahead, or reads from zero if the
    seek is behind the current position."""
    curr_offset = self.tell()
    new_offset = offset # Default to seek from start
    if from_what == 1: # Seek from current offset
        new_offset = curr_offset + offset
    elif from_what == 2: # Seek from EOF
        new_offset = self._seeklen + offset

    if new_offset > self._seeklen:
        new_offset = self._seeklen

    if new_offset < 0:
        new_offset = 0

    if new_offset > curr_offset:
        # If the new position is after the current,
        # then read up to the new position
        self.read(new_offset - curr_offset)
    elif new_offset < curr_offset:
        if self._offset >= curr_offset - new_offset:
            self._offset -= curr_offset - new_offset
        else:
            # Position is before the current position. Reset the ZipExtFile
            # object and read up to the new position.
            self._fileobj.seek(self._startcomp)
            self._running_crc = self._startcrc
            self._compress_left = self._complen
            self._left = self._seeklen
            self._readbuffer = b''
            self._offset = 0
            self._decompressor = zipfile._get_decompressor(self._compress_type)
            self._eof = False
            # If the new position is the start of file, don't read
            if new_offset > 0:
                self.read(new_offset)

    return self.tell()

def tellZipExtFile(self):
    offset = self._seeklen - self._left - len(self._readbuffer) + self._offset
    return offset

def makeseekable(zh):
    # Keep track of the start of compressed data in order to reset to that point
    zh._startcomp = zh._fileobj.tell()
    # This is usually zero, but no chances were taken
    zh._startcrc  = zh._running_crc
    # These two values need to be restored when the object is reset
    zh._complen   = zh._compress_left
    zh._seeklen   = zh._left
    # Yes, yes we are
    zh.seekable   = lambda: True
    zh.seek       = types.MethodType(seekZipExtFile, zh)
    zh.tell       = types.MethodType(tellZipExtFile, zh)
    return zh

def fileinfo(fh):
    fs = 0
    fc = 0
    fh.seek(0)
    for piece in fh:
        fc = zlib.crc32(piece, fc)
        fs += len(piece)

    return fs, fc

def filestat(fh):
    if zipfile.is_zipfile(fh):
        print("{} is a zipfile".format(fh.name))
        with zipfile.ZipFile(fh) as zf:
            for zi in zf.infolist():
                with zf.open(zi) as zh:
                    filestat(makeseekable(zh))
    else:
        fsize, fcrc = fileinfo(fh)
        print("Size: {}, CRC: 0x{:08X}".format(fsize, fcrc))

def dirstat(fp):
    print(fp)
    if os.path.isdir(fp):
        for fname in os.listdir(fp):
            dirstat(os.path.join(fp, fname))
    elif os.path.isfile(fp):
        filestat(open(fp, 'rb'))
    else:
        print("I can't process {}".format(fp))

if __name__ == '__main__':
    aparse = argparse.ArgumentParser(description="Zip-in-zip test")
    aparse.add_argument("zipfile")

    args = aparse.parse_args()

    dirstat(args.zipfile)