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




class PyinstHeader(Struct):

    def parse(self):
        yield String(8, zero_terminated=True, name="Signature")
        yield UInt32BE(name="PackageSize", comment="size of package")
        yield UInt32BE(name="TableOfContentOffset", comment="toc offset")
        yield UInt32BE(name="TableOfContentSize", comment="toc size")
        yield UInt32BE(name="PythonVersion", comment="python version")


class PyinstEntry(Struct):

    def parse(self):
        sz = yield UInt32BE(name="Size", comment="size of entry")
        yield UInt32BE(name="Offset", comment="offset of entry")
        yield UInt32BE(name="CompressedSize", comment="size of entry compressed")
        yield UInt32BE(name="UncompressedSize", comment="size of entry uncompressed")
        yield UInt8(name="CompressionFlags", comment="")
        yield UInt8(name="Type", comment="type of entry")
        rest = sz - len(self)
        if rest > 0:
            yield String(rest, name="Name", zero_terminated=True, comment="entry name")
        

class PyinstAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "PYINST"
    regexp = r"MEI\x0C\x0B\n\x0B\x0E"

    @classmethod
    def locate(cls, curfile, offset_magic, parent_parser):
        if parent_parser is None and curfile[0:1] == b"\x78":
            return 0, ""
        elif parent_parser is not None and parent_parser.name == "PE":
            max_fend = 0
            for section in parent_parser["Sections"]:
                # pyinst exe: archive should start at overlay
                max_fend = max(max_fend, section["PointerToRawData"] + section["SizeOfRawData"])
            if max_fend < offset_magic:
                return max_fend, ""
        elif parent_parser is not None and parent_parser.name == "ELF":
            sec_start = None
            for section in parent_parser["Sections"]:
                if section["Offset"] <= offset_magic and section["Offset"] + section["Size"] > offset_magic:
                    # pyinst exe: archive should be in a .pydata section
                    sec_start = section["Offset"]
                    break
            if sec_start is not None:
                return sec_start, ""
        return None

    def __init__(self):
        FileTypeAnalyzer.__init__(self)
        self.filesystem = {}

    def open(self, vfile, password=None):
        self.set_architecture(malcat.Architecture.PY36)
        entry = self.filesystem[vfile.path]
        data = self.read(entry["Offset"], entry["CompressedSize"])
        if entry["CompressedSize"] == entry["UncompressedSize"] or entry["CompressionFlags"] == 0:
            pass
        else:
            import zlib
            data = zlib.decompress(data)
        if entry["Type"] in (115, 109):
            # we need to prepend python header
            pyversion = self["PyinstHeader"]["PythonVersion"]
            if pyversion > 100:
                pyversion = (pyversion//100, pyversion % 100)
            else:
                pyversion = (pyversion//10, pyversion % 10)
            MAGIC_TO_VERSION_INV = {v: k for k, v in MAGIC_TO_VERSION.items()}
            magic = MAGIC_TO_VERSION_INV.get(pyversion, MAGIC_TO_VERSION_INV[(3, 6)])
            header = struct.pack("<I", magic)
            header += 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):
        magic_off, _ = self.search(r"MEI\x0C\x0B\n\x0B\x0E")
        if magic_off is None:
            raise FatalError

        # pyinst
        self.jump(magic_off)
        pyinst = yield PyinstHeader(category=Type.HEADER)
        # pyinst table of content
        toc = pyinst["TableOfContentOffset"]
        toc_end = toc + pyinst["TableOfContentSize"]
        self.jump(toc)
        while self.tell() < toc_end:
            e = yield PyinstEntry(parent=pyinst)
            ename = "Name" in e and e["Name"].replace("\0", " ").strip() or ""
            self.add_section(ename, e["Offset"], e["CompressedSize"])
            if ename:
                ename = ename.replace("\\", "/")
                self.filesystem[ename] = e
                self.add_file(ename, e["UncompressedSize"], "open")
            if e["Offset"] + e["CompressedSize"] > self.size(): 
                raise OutOfBoundError
            if e["Type"] == ord('z') and self.read(e["Offset"], 3) != b"PYZ":
                raise FatalError
            if e["CompressionFlags"] == 1 and self.read(e["Offset"], 1) != b"\x78":
                raise FatalError
            self.confirm()

            
