from filetypes.base import *
import malcat

BUFFER_INCREMENT = 1024*16

class StreamHeader(Struct):

    def parse(self):
        yield UInt8(name="CMF", comment="compression method")
        yield UInt8(name="FLG", comment="flags")

class ZlibAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "ZLIB"
    regexp = r"\x78[\x9C\xDA]"

    def unpack(self, vfile, password=None):
        import zlib
        stream = self.at("Deflate stream")
        return zlib.decompress(self.read(0, stream.offset + stream.size))

    def parse(self, hint):
        import zlib
        ptr = self.tell()
        remaining = self.remaining()
        yield StreamHeader(category=Type.HEADER)
        # bruteforce deflate stream length
        obj = zlib.decompressobj()
        decompressed_size = 0
        while remaining and not obj.eof:
            todo = min(remaining, BUFFER_INCREMENT)
            try:
                decompressed_size += len(obj.decompress(self.read(ptr, todo)))
            except BaseException as e:
                raise FatalError("Not a valid Zlib stream ({})".format(e))
            ptr += todo
            remaining -= todo
        if not obj.eof:
            raise FatalError("Truncated Zlib stream")
        packed_size = ptr - (2 + len(obj.unused_data))
        self.add_section("stream", self.tell(), packed_size)
        yield Bytes(packed_size, name="Deflate stream", category=Type.DATA)
        self.add_file("<packed content>", decompressed_size, "unpack")


