from filetypes.base import *
import malcat

BUFFER_INCREMENT = 1024*16

class BZHeader(Struct):

    def parse(self):
        yield String(3, name="Signature")
        yield String(1, zero_terminated=False, name="CompressionLevel", comment="'1'..'9' block-size 100 kB-900 kB (uncompressed)")
        yield String(6, name="Magic", comment="0x314159265359 (BCD (pi))")


class BZ2Analyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "BZ2"
    regexp = r"BZh[1-9]1AY&SY"


    def unpack(self, vfile, password=None):
        import bz2
        return bz2.decompress(self.read(0, vfile.size))
        
    def parse(self, hint):
        import bz2
        ptr = self.tell()
        remaining = self.remaining()
        h = yield BZHeader(category=Type.HEADER)
        # bruteforce deflate stream length
        obj = bz2.BZ2Decompressor()
        while remaining and not obj.eof:
            todo = min(remaining, BUFFER_INCREMENT)
            try:
                obj.decompress(self.read(ptr, todo))
            except BaseException as e:
                raise FatalError("Not a valid Bz2 stream ({})".format(e))
            ptr += todo
            remaining -= todo
        if not obj.eof:
            raise FatalError("Truncated Bz2 stream")
        yield Bytes(ptr - (10 + len(obj.unused_data)), name="BZip2 stream", category=Type.DATA)
        self.add_file("<content>", self.tell(), "unpack")

        
