from filetypes.base import *
import malcat
import struct
import re
import zlib
from malcat import PythonHelper


class InstallShieldHeader(Struct):

    def parse(self):
        yield String(14, zero_terminated=True, name="Magic")
        yield UInt16(name="NumberOfFiles")
        yield UInt32(name="Version")
        yield Bytes(8, name="x4")
        yield UInt16(name="x5")
        yield Unused(16)


class ISFileAttribute4(Struct):

    def parse(self):
        sn = yield UInt32(name="FileNameLength")
        if sn == 0 or sn > 1024:
            raise FatalError("Invalid file name size")
        yield UInt32(name="EncodedFlags")        
        yield UInt16(name="x3")        
        yield UInt32(name="FileSize")        
        yield Unused(8)
        yield UInt16(name="IsUnicodeLauncher")
        yield Unused(24)
        yield StringUtf16le(sn // 2, name="FileName")


class ISFileAttribute3(Struct):

    def parse(self):
        sn = yield UInt32(name="FileNameLength")
        if sn == 0 or sn > 1024:
            raise FatalError("Invalid file name size")
        yield UInt32(name="EncodedFlags")        
        yield UInt16(name="x3")        
        yield UInt32(name="FileSize")        
        yield Unused(8)
        yield UInt16(name="IsUnicodeLauncher")
        yield StringUtf16le(sn // 2, name="FileName")        



class InstallShieldAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "InstallShield"
    regexp = r"ISSetupStream\x00.{32}..\x00\x00\x06\x00\x00\x00"


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

    #def decode_data_block(self, data, key, key_offset=0):
    #    data = memoryview(data)
    #    res = bytearray(len(data))
    #    decoded = 0
    #    while decoded < len(res):
    #        decode_start = (decoded + key_offset) % 1024;
    #        todo = min(1024 - decode_start, len(data) - decoded)
    #        for i, c in enumerate(data[decoded:decoded + todo]):
    #            res[decoded + i] = ~((c << 4 | c >> 4) ^ key[(i+decode_start) % len(key)]) & 0xff
    #        decoded += todo
    #    return res

    def open(self, vfile, password=None):
        if not vfile.path in self.filesystem:
            raise KeyError(f"File not found: {vfile.path}")
        ofname, offset, size, flags, needs_inflate = self.filesystem[vfile.path]
        encoded_block_len = 0x4000
        key = bytearray(ofname.encode("utf8"))
        if flags & 4:
            encoded_block_len = 1024
        len_read = min(4096 * 64, encoded_block_len)
        done = 0
        res = []
        while done < size:
            len_read = min(len_read, size - done)
            data = self.read(offset + done, len_read)
            decoded_data = PythonHelper.installshield_decrypt(data, key, 0)
            done += len(data)
            res.append(decoded_data)
        if (flags & 6) == 2:
            pass
        res = b"".join(res)
        if needs_inflate:
            res = zlib.decompress(res)
        return res

    def parse(self, hint):
        hdr = yield InstallShieldHeader()
        if hdr["Version"] == 4:
            fa_class = ISFileAttribute4
        elif hdr["Version"] == 3:
            fa_class = ISFileAttribute3
        else:
            raise FatalError(f"Unsupported installshield type {hdr['Type']}")

        for i in range(hdr["NumberOfFiles"]):
            start = self.tell()
            fa = yield fa_class(name="FileAttribute")
            filesize = fa["FileSize"]
            fname = fa["FileName"].replace("\\", "/")
            self.filesystem[fname] = (fa["FileName"], self.tell(), filesize, fa["EncodedFlags"], fa["IsUnicodeLauncher"])
            self.jump(self.tell() + filesize)
            self.add_section(fname, start, self.tell() - start)
            self.add_file(fname, filesize, "open")
            self.confirm()
