from filetypes.base import *
from filetypes.BMP import RGBA
import malcat
import struct
import re

# see https://github.com/kichik/nsis/blob/master/Source/exehead/fileform.h
def nsis_lzma(mv):
    import lzma
    lzma_filter = lzma._decode_filter_properties(lzma.FILTER_LZMA1, mv[0:5])
    return lzma.decompress(mv[5:], lzma.FORMAT_RAW, filters=[lzma_filter])

def nsis_bz2(mv):
    if len(mv) < 2 or mv[0] != 0x31 or mv[1] >= 0xe:
        raise ValueError("Not a valid bz2 stream")
    decompressor = malcat.NsisBzip2Decompressor()
    return decompressor.decompress(mv)


class FirstHeader(Struct):

    def parse(self):
        yield BitsField(
             Bit(name="Uninstall", comment="this is an NSIS uninstaller"),
             Bit(name="Silent", comment="installation is silent"),
             Bit(name="NoCRC", comment=""),
             Bit(name="ForceCRC", comment=""),
             NullBits(28),
            name="Flags")
        yield UInt32(name="Signature", comment="should be 0xDeadBeef unless patched")
        yield String(12, zero_terminated=False, name="Magic", comment="should be NullsoftInst")
        yield UInt32(name="InstallerSize", comment="")
        yield UInt32(name="ArchiveSize", comment="")


class NSISAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "NSIS"
    regexp = r"(?<=[\x00-\x0f]\x00\x00\x00)\xef\xbe\xad\xdeNullsoftInst"

    @classmethod
    def locate(cls, curfile, offset_magic, parent_parser):
        return offset_magic - 4, ""

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

    def parse(self, hint):
        hdr = yield FirstHeader(category=Type.HEADER)
        if hdr["ArchiveSize"] > self.size() or hdr["ArchiveSize"] < hdr.size:
            raise FatalError("Invalid archive size")
        self.add_section("PackedData", hdr.offset + hdr.size, hdr["ArchiveSize"] - hdr.size, discardable=True)
        self.confirm()
        self.set_eof(hdr["ArchiveSize"])
        try:
            archive = self.decompress()
        except NotImplementedError as e:
            raise FatalError(e)
        except BaseException as e:
            raise FatalError(e)
        if not archive:
            raise FatalError("Empty NSIS")

        # try to decompress SETUP, because archive file names can only be obtained by disassembling NSIS code
        files_offsets = {}
        try:
            setup_code = b""
            strings_start = None
            setup_analyzer = NSISSetupAnalyzer()
            setup_file = malcat.FileBuffer(bytearray(archive[0][0]), "SETUP")
            setup_analyzer.run(setup_file)
            setup_code = setup_file[setup_analyzer.code_start : setup_analyzer.code_start + setup_analyzer.code_size]
            strings_start = setup_analyzer.string_start

            if strings_start is None:
                raise ValueError("no string section")
            is_unicode = setup_file[strings_start:strings_start+2] == b"\x00\x00"
            # we don't have access to the disassembler yet, so just use re to look for ExtractFile opcode
            for m in re.finditer(rb"\x14\x00\x00\x00(?:[\x90-\x9f]\x00\x00\x05|.\x00\x00[\x20\x02])(....)(....)", setup_code, flags=re.DOTALL):
                string_offset, = struct.unpack("<I", m.group(1))
                offset, = struct.unpack("<I", m.group(2))
                if is_unicode:
                    string = setup_file.read_cstring_utf16le(strings_start + string_offset * 2, 256)
                else:
                    string = setup_file.read_cstring_ascii(strings_start + string_offset, 256)
                if string:
                    files_offsets[offset] = string
        except BaseException as e:
            print(e)
        self.filesystem = {
            "#Setup": 0
        }
        for i in range(1, len(archive)):
            offset = archive[i][1]
            name = files_offsets.get(offset, f"FILE{i}-#0x{offset:x}")
            self.filesystem[name] = i

        for name, index in self.filesystem.items():
            hint = ""
            if name == "#Setup":
                hint = "NSIS.Setup"
            self.add_file(name, len(archive[index][0]), "unpack", hint)


    def decompress_solid(self, data):
        import zlib
        res = []
        hdr = self["FirstHeader"]
        #try lzma
        unpacked = None
        try:
            unpacked = nsis_lzma(data[hdr.size:])
        except Exception: pass
        #try zlib
        if unpacked is None:
            try:
                unpacked = zlib.decompress(data[hdr.size:], -zlib.MAX_WBITS)
            except Exception: pass
        #try bz2
        if unpacked is None:
            try:
                unpacked = nsis_bz2(data[hdr.size:])
            except NotImplementedError: raise
            except Exception: raise
        if unpacked is None:
            return res
        p = 0
        while p + 4 < len(unpacked):
            fsize, = struct.unpack("<I", unpacked[p:p+4])
            if not res:
                res.append((unpacked[p+4:p+4+fsize], None))
            else:
                res.append((unpacked[p+4:p+4+fsize], p - (len(res[0][0]) + 4)))
            p = p + 4 + fsize
        return res

    def decompress(self):
        import zlib
        res = []
        hdr = self["FirstHeader"]
        data = self.read(0, hdr["ArchiveSize"])
        installer_size_packed, = struct.unpack("<I", data[hdr.size: hdr.size + 4])
        installer_size_packed = installer_size_packed & 0x7fffffff
        installer_packed_data = data[hdr.size + 4: hdr.size + 4 + installer_size_packed]
        installer_data = None
        ## try to decrypt header data
        is_lzma = False
        is_zlib = False
        is_bz2 = False
        if installer_size_packed == hdr["InstallerSize"]:
            installer_data = installer_packed_data
        #try lzma
        if installer_data is None:
            try:
                installer_data = nsis_lzma(installer_packed_data)
                is_lzma = True
            except Exception: pass
        #try zlib
        if installer_data is None:
            try:
                installer_data = zlib.decompress(installer_packed_data, -zlib.MAX_WBITS)
                is_zlib = True
            except Exception: pass
        #try bz2
        if installer_data is None:
            try:
                installer_data = nsis_bz2(installer_packed_data)
                is_bz2 = True
            except NotImplementedError: raise
            except Exception: pass            
        if installer_data is None:
            res = self.decompress_solid(data)
            if not res:
                raise FatalError("Could not figure out which version/encryption is used !")
            self.is_solid = True
            return res
        self.is_solid = False
        res.append((installer_data, None))
        p = hdr.size + 4 + installer_size_packed
        while p + 4 < len(data):
            size_packed, = struct.unpack("<I", data[p: p + 4])
            compressed = (size_packed & 0x80000000) != 0
            size_packed = size_packed & 0x7fffffff
            try:
                if not compressed:
                    unpacked = data[p+4:p+4+size_packed]
                elif is_lzma:
                    unpacked = nsis_lzma(data[p+4:p+4+size_packed])
                elif is_zlib:
                    unpacked = zlib.decompress(data[p+4:p+4+size_packed], -15)
                elif is_bz2:
                    unpacked = nsis_bz2(data[p+4:p+4+size_packed])
                else:
                    unpacked = data[p+4:p+4+size_packed]
            except Exception as e:
                print(e)
                unpacked = str(e).encode("utf8")
                unpacked = data[p+4:p+4+size_packed]
            res.append((unpacked, p - (hdr.size + 4 + installer_size_packed)))
            p = p + 4 + size_packed
        return res
        

    def unpack(self, vfile, password=None):
        archive_index = self.filesystem[vfile.path]
        return self.decompress()[archive_index][0]


        
        
        


BLOCKS_NUM = 8
NSIS_MAX_INST_TYPES = 32


class BlockHeader(Struct):

    def parse(self):
        yield Va32(name="Offset")
        yield UInt32(name="Num")

class SectionHeader(Struct):

    def __init__(self, is_unicode, *args, **kwargs):
        Struct.__init__(self, *args, **kwargs)
        self.is_unicode = is_unicode

    def parse(self):
        yield Va32(name="NamePointer")
        yield UInt32(name="InstallTypes")
        yield BitsField(
             Bit(name="Selected"),
             Bit(name="SectionGroup"),
             Bit(name="SectionGroupEnd"),
             Bit(name="Bold"),
             Bit(name="ReadOnly"),
             Bit(name="Expand"),
             Bit(name="PSelected"),
             Bit(name="Toggled"),
             Bit(name="NameChange"),
             NullBits(23),
            name="Flags")
        yield UInt32(name="CodeIndex")
        yield UInt32(name="CodeSize")
        yield UInt32(name="SizeKilobytes")
        if self.is_unicode:
            yield String(2048, name="Name", zero_terminated=True)
        else:
            yield String(1024, name="Name", zero_terminated=True)


NSIS_CALLBACKS = (
        "OnInit",
        "OnInstSuccess",
        "OnInstFailed",
        "OnUserAbort",
        "OnGuiInit",
        "OnGuiEnd",
        "OnMouseOverSection",
        "OnVerifyInstDir",
        "OnSelChange",
        "OnRebootFailed",
        )

class CommonHeader(Struct):

    def parse(self):
        yield BitsField(
             Bit(name="ShowDetails"),
             Bit(name="NeverShow", comment=""),
             Bit(name="ProgressColor", comment=""),
             Bit(name="Silent", comment=""),
             Bit(name="SilentLog", comment=""),
             Bit(name="AutoClose", comment=""),
             Bit(name="DirNoShow", comment=""),
             Bit(name="NoRootDir", comment=""),
             Bit(name="CompCustom", comment=""),
             Bit(name="NoCustom", comment=""),
             NullBits(22),
            name="Flags")
        yield Array(BLOCKS_NUM, BlockHeader(), name="Blocks")
        yield UInt32(name="InstallRegRootkey", comment="InstallDirRegKey stuff")
        yield UInt32(name="InstallRegKeyPtr", comment="ignored")
        yield UInt32(name="InstallRegValuePtr", comment="ignored")
        yield RGBA(name="BackgroundColor1", comment="#ifdef NSIS_SUPPORT_BGBG")
        yield RGBA(name="BackgroundColor2", comment="#ifdef NSIS_SUPPORT_BGBG")
        yield RGBA(name="TextColor", comment="#ifdef NSIS_SUPPORT_BGBG")
        yield RGBA(name="LogBackroundColor", comment="installation log window colors")
        yield RGBA(name="LogForegroundColor", comment="installation log window colors")
        yield UInt32(name="LangTableSize", comment="langtable size")
        yield RGBA(name="LicenseBackgroundColor", comment="license background color")
        for cb in NSIS_CALLBACKS:
            yield UInt32(name="Callback{}".format(cb), comment="callback")
        yield Array(NSIS_MAX_INST_TYPES + 1, UInt32(), name="InstallTypes")
        yield UInt32(name="InstallDirPtr", comment="default install dir")
        yield UInt32(name="InstallDirAutoAppend", comment="auto append part")
        yield UInt32(name="StringUninstallChild", comment="")
        yield UInt32(name="StringUninstallCmd", comment="")
        yield UInt32(name="StringWinInit", comment=" path of wininit.ini")



class NSISSetupAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "NSIS.Setup"


    def __init__(self):
        FileTypeAnalyzer.__init__(self)
        self.code_start = 0
        self.string_start = 0
        self.nsis_sections = {}

    def read_string(self, idx):
        if idx == 0xffffffff:
            return ""
        if self.is_unicode:
            return self.read_cstring_utf16le(self.string_start + idx*2)
        else:
            return self.read_cstring_ascii(self.string_start + idx)

    def parse(self, hint):
        self.set_architecture(malcat.Architecture.NSIS)
        hdr = yield CommonHeader(category=Type.HEADER)
        self.confirm()
        self.code_start = hdr["Blocks"][2]["Offset"]
        self.string_start = hdr["Blocks"][3]["Offset"]
        self.code_size = self.string_start - self.code_start
        self.is_unicode = self.read(self.string_start, 2) == b"\x00\x00"
        # parse sections
        sec_off = hdr["Blocks"][1]["Offset"]
        self.jump(sec_off)
        sections = yield Array(hdr["Blocks"][1]["Num"], SectionHeader(self.is_unicode), name="Sections", category=Type.HEADER)
        for s in sections:
            if s["NamePointer"]:
                name = self.read_string(s["NamePointer"])
            else:
                name = s["Name"].replace("\x00", "")
            self.nsis_sections[name] = s
            if s["CodeSize"]:
                self.add_symbol(self.code_start + s["CodeIndex"] * 4 * 7, name or "EntryPoint", malcat.FileSymbol.ENTRY)
        blocknames = ["Pages", "Sections", "Entries", "Strings", "LangTable", "CtlColors", "BgFont", "Data"]
        for i, name in enumerate(blocknames):
            block = hdr["Blocks"][i]
            if i < len(blocknames) - 1 and hdr["Blocks"][i+1]["Offset"]:
                sz = hdr["Blocks"][i+1]["Offset"] - block["Offset"]
            else:
                sz = self.size() - block["Offset"] 
            if block["Offset"] and sz > 0:
                self.add_section(name, block["Offset"], sz, x=name=="Entries")
               

        # parse symbols
        for cb in NSIS_CALLBACKS:
            entry = hdr["Callback{}".format(cb)]
            if entry != 0xffffffff:
                self.add_symbol(self.code_start + entry * 4 * 7, cb, malcat.FileSymbol.ENTRY)
                

