from filetypes.base import *
from filetypes.PYC import *
import malcat
import struct


class PyzHeader(Struct):

    def parse(self):
        yield String(4, zero_terminated=True, name="Signature")
        magic = yield UInt32(name="PythonMagic", comment="python magic value")
        version = MAGIC_TO_VERSION.get(magic, None)
        if version is None:
            version = MAGIC_TO_VERSION.get(magic + 1, None)
        if version is None:
            raise FatalError("Bad magic")
        yield Int32BE(name="TableOfContent", comment="offset of archve table of content")


        

class PyzAnalyzer(PythonAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "PYZ"
    regexp = r"PYZ\x00"


    def __init__(self):
        PythonAnalyzer.__init__(self)

    
    def open(self, vfile, password=None):
        import zlib
        toc = self.read_object(self["TableOfContent"])
        if type(toc) == list:
            toc = dict(toc)
        package, offset, size = toc[vfile.path]
        data = self.read(offset, size)
        data = zlib.decompress(data)
        # append python header
        magic = self["PyzHeader"]["PythonMagic"]
        pyversion = MAGIC_TO_VERSION[magic]
        header = struct.pack("<I", magic) + b"\x00\x00\x00\x00"
        if pyversion >= (3, 7):
            header += b"\x00\x00\x00\x00"
        if pyversion >= (3, 3):
            header += struct.pack("<I", len(data))
        data = header + data
        return bytearray(data)

    def parse(self, hint):
        self.set_architecture(malcat.Architecture.PY36)
        pyz = yield PyzHeader(category=Type.HEADER)
        self.confirm()
        # pyz table of content
        toc_address = pyz["TableOfContent"]
        self.jump(toc_address)
        typ, = struct.unpack("<B", self.read(self.tell(), 1))
        typ = typ & 0x7f
        if typ == TYPE_LIST:
            toc = yield PySequence(name="TableOfContent",category=Type.HEADER)
        else:
            toc = yield PyDict(name="TableOfContent", category=Type.HEADER)

        # parse toc
        toc = self.read_object(toc)
        if type(toc) == list:
            toc = dict(toc)
        for name, infos in toc.items():
            package, offset, size = infos
            if type(package) != str:
                package = name
            if type(offset) != int or type(size) != int or type(package) != str:
                print("Invalid TOC entry: {}, {}, {}".format(package, offset, size))
                continue
            self.add_section(name, offset, size)
            self.add_file(name, size, "open")
