from filetypes.base import *
import malcat

BUFFER_INCREMENT = 1024*16

class GzipHeader(Struct):

    def parse(self):
        yield UInt16(name="Signature", comment="always 1F8B")
        yield UInt8(name="Compression", comment="compression method used", values=[
            ("store (copied)", 0x00),
            ("compress", 0x01),
            ("pack", 0x02),
            ("lzh", 0x03),
            ("deflate", 0x08),
            ])
        flags = yield BitsField(
                Bit(name="FileIsAscii", comment="file probably contains ascii text"),
                Bit(name="FileIsMultiPart", comment="continuation of multi-part gzip file, part number present"),
                Bit(name="FileHasExtraField", comment="extra field present"),
                Bit(name="FileHasName", comment="original file name present"),
                Bit(name="FileHasComment", comment="file comment present"),
                Bit(name="FileIsEncrypted", comment="file is encrypted, encryption header present"),
                name="Flags")
        yield Timestamp(name="ModificationTime", comment="file modification time")
        yield UInt8(name="ExtraFlags", comment="extra flags (depend on compression method)")
        yield UInt8(name="OperatingSystem", comment="OS type", values=[
            ("FAT filesystem (MS-DOS, OS/2, NT/Win32)", 0x00),
            ("Amiga", 0x01),
            ("VMS (or OpenVMS)", 0x02),
            ("Unix", 0x03),
            ("VM/CMS", 0x04),
            ("Atari TOS", 0x05),
            ("HPFS filesystem (OS/2, NT)", 0x06),
            ("Macintosh", 0x07),
            ("Z-System", 0x08),
            ("CP/M", 0x09),
            ("TOPS-20", 0x0a),
            ("NTFS filesystem (NT)", 0x0b),
            ("QDOS", 0x0c),
            ("Acorn RISCOS", 0x0d),
            ("unknown", 0xff),
            ])
        if flags["FileIsMultiPart"]:
            yield UInt16(name="PartNumber", comment="optional part number (second part=1)")
        if flags["FileHasExtraField"]:
            l = yield UInt16(name="ExtraFieldLength", comment="extra field length")
            yield Bytes(l, name="ExtraField", comment="extra field")
        if flags["FileHasName"]:
            yield CString(name="FileName", comment="original file name, zero terminated")
        if flags["FileHasComment"]:
            yield CString(name="Comment", comment="file comment, zero terminated")
        if flags["FileIsEncrypted"]:
            yield Bytes(12, name="EncryptionHeader", comment="encryption header")

class GzipFooter(Struct):

    def parse(self):
        yield UInt32(name="Crc32", comment="crc32")
        yield UInt32(name="UncompressedInputSize", comment="uncompressed input size modulo 2^32")


class GzipAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "GZIP"
    regexp = r"\x1f\x8b[\x00-\x08][\x00-\x3f].....[\x00-\x0d\xff]"

    def unpack(self, vfile, password=None):
        import zlib
        footer = self["GzipFooter"]
        size = footer.offset + footer.size
        obj = zlib.decompressobj(wbits=31)
        return bytearray(obj.decompress(self.read(0, size)))

    def parse(self, hint):
        import zlib
        ptr = self.tell()
        remaining = self.remaining()
        hdr = yield GzipHeader(category=Type.HEADER)
        data_pos = self.tell()
        hdr_size = data_pos - ptr
        # bruteforce stream length
        obj = zlib.decompressobj(wbits=31)
        self.jump(ptr)
        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 GZip stream ({})".format(e))
            ptr += todo
            remaining -= todo
        if not obj.eof:
            raise FatalError("Truncated GZip stream")
        self.jump(data_pos)
        yield Bytes(ptr - (hdr_size + len(obj.unused_data) + 8), name="Stream", category=Type.DATA)
        footer = yield GzipFooter(category=Type.HEADER)
        self.add_file("FileName" in hdr and hdr["FileName"] or "<content>", footer["UncompressedInputSize"], "unpack")
