from filetypes.base import *
import malcat
import struct
from transforms.substitution import VbeDecrypt


class VbeHeader(Struct):

    def parse(self):
        yield String(4, name="Magic")
        yield String(8, name="SizePacked", comment="base64-encoded size of packed content")

    
class VbeFooter(Struct):

    def parse(self):
        yield String(8, name="SizeUnpacked", comment="base64-encoded size of unpacked content")
        yield String(4, name="Magic")

class VbeAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.PROGRAM
    name = "VBE"
    regexp = r"#@~\^......=="

    def unpack(self, vfile, password=None):
        data = self["Content"]
        return VbeDecrypt().run(data)

    def decompile(self):
        data = self["Content"]
        return VbeDecrypt().run(data).decode("latin1")

    def parse(self, hint):
        import base64
        hdr = yield VbeHeader()
        self.confirm()
        sz, = struct.unpack("<I", base64.b64decode(hdr["SizePacked"]))
        start = self.tell()
        yield Bytes(sz, name="Content", category=Type.DATA)
        end = self.tell()
        footer = yield VbeFooter()
        self.add_file("<content>", sz, "unpack")

        self.add_section("Content", start, end-start)
        # sometimes it's still valid
        if footer["Magic"] != "^#~@":
            raise FatalError("Invalid footer")

  
