from filetypes.base import *
import malcat
import struct



class CFHeader(Struct):

    def parse(self):
        yield String(4, name="Signature")
        yield UInt32(name="Reserved", comment="should be 0")
        yield UInt32(name="Size", comment="cabinet file size")
        yield UInt32(name="Reserved2", comment="should be 0")
        yield Offset32(name="FirstFile", comment="offset of first cfffile entry")
        yield UInt32(name="Reserved3", comment="should be 0")
        yield UInt8(name="VersionMinor", comment="should be 3")
        yield UInt8(name="VersionMajor", comment="should be 1")
        yield UInt16(name="NumberOfFolders", comment="number of folders")
        yield UInt16(name="NumberOfFiles", comment="number of files")
        flags = yield BitsField(
            Bit(name="PrevCabinet", comment="The flag is set if this cabinet file is not the first in a set of cabinet files. When this bit is set, the szCabinetPrev and szDiskPrev fields are present in this CFHEADER structure"),
            Bit(name="NextCabinet", comment=" The flag is set if this cabinet file is not the last in a set of cabinet files. When this bit is set, the szCabinetNext and szDiskNext fields are present in this CFHEADER structure"),
            Bit(name="ReservePresent", comment="The flag is set if if this cabinet file contains any reserved fields. When this bit is set, the cbCFHeader, cbCFFolder, and cbCFData fields are present in this CFHEADER structure"),
            NullBits(13),
            name="Flags", comment="bit-mapped values that indicate the presence of optional data")
        yield UInt16(name="SetId", comment="Specifies an arbitrarily derived (random) value that binds a collection of linked cabinet files together. All cabinet files in a set will contain the same setID field value")
        yield UInt16(name="iCabinet", comment="Specifies the sequential number of this cabinet in a multicabinet set. The first cabinet has iCabinet=0")
        if flags["ReservePresent"]:
            reserve = yield UInt16(name="ReserveHeader", comment="size, in bytes, of the abReserve field in this CFHEADER structure")
            yield UInt8(name="ReserveFolder", comment="size, in bytes, of the abReserve field in each CFFOLDER field entry")
            yield UInt8(name="ReserveData", comment="size, in bytes, of the abReserve field in each CFDATA field entry")
            yield Bytes(reserve, name="ReservedData")

class CFFolder(Struct):

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

    def parse(self):
        yield Offset32(name="DataStart", comment="offset of the first CFDATA block in this folder")
        yield UInt16(name="DataCount", comment="number of CFDATA blocks in this folder")
        yield UInt8(name="CompressionMethod", comment="compression type indicator", values=[
            ("None", 0),
            ("MSZIP", 1),
            ("Quantum", 2),
            ("LZX", 3),
            ])
        yield UInt8(name="CompressionLevel", comment="compression level (for lzx and quantum)")
        if self.reserved:
            yield Bytes(self.reserved, name="Reserved")


class CFFile(Struct):

    def parse(self):
        yield UInt32(name="UncompressedSize", comment="uncompressed size of this file in bytes")
        yield UInt32(name="OffsetFolderStart", comment="uncompressed byte offset of the start of this file's data. For the first file in each folder, this value will usually be zero. Subsequent files in the folder will have offsets that are typically the running sum of the cbFile values.")
        yield UInt16(name="FolderIndex", comment="index of the folder containing this file's data. A value of zero indicates this is the first folder in this cabinet file. The special iFolder values ifoldCONTINUED_FROM_PREV and ifoldCONTINUED_PREV_AND_NEXT indicate that the folder index is actually zero, but that extraction of this file would have to begin with the cabinet named in CFHEADER.szCabinetPrev. The special iFolder values ifoldCONTINUED_PREV_AND_NEXT and ifoldCONTINUED_TO_NEXT indicate that the folder index is actually one less than CFHEADER.cFolders, and that extraction of this file will require continuation to the cabinet named in CFHEADER.szCabinetNext")
        yield DosDate(name="Date", comment="date stamp for this file")
        yield DosTime(name="Time", comment="time stamp for this file")
        attr = yield BitsField(
            Bit(name="ReadOnly"),
            Bit(name="Hidden"),
            Bit(name="System"),
            NullBits(2),
            Bit(name="Archive"),
            Bit(name="Executable", comment="run after extraction"),
            Bit(name="UtfName", comment="szName[] contains UTF "),
            NullBits(8),
            name="Attributes", comment="file attributes")
        if attr["UtfName"]:
            yield CStringUtf8(name="FileName")
        else:
            yield CString(name="FileName")


class CFData(Struct):

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

    def parse(self):
        yield UInt32(name="Checksum", comment="checksum of this entry")
        cs = yield UInt16(name="CompressedSize", comment="number of bytes in this block")
        yield UInt16(name="UncompressedSize", comment="number of bytes after decompression")
        if self.reserved:
            yield Bytes(self.reserved, name="Reserved")
        yield Bytes(cs, name="Data")


class Decompressor:

    def decompress_one_block(self, data, decompressed_size):
        if len(data) > decompressed_size:
            raise ValueError("decompressed size {} vs {}".format(len(data), decompressed_size))
        return data

class ZlibDecompressor:

    def __init__(self):
        self.zdict = b""

    def decompress_one_block(self, data, decompressed_size):
        import zlib
        #compressed = decompressed_size == len(data)
        data = data[2:]
        #if not compressed and False:
        #    return data
        decompress = zlib.decompressobj(-zlib.MAX_WBITS, zdict=self.zdict)
        res = decompress.decompress(data)
        res += decompress.flush()
        if len(res) != decompressed_size:
            raise ValueError("decompressed size {} vs {}".format(len(res), decompressed_size))
        self.zdict = res
        return res

class LzxDecompressor:

    def __init__(self, window_size):
        self.dec = malcat.LzxDecompressor(window_size)

    def decompress_one_block(self, data, decompressed_size):
        res = self.dec.decompress(data, decompressed_size)
        if len(res) != decompressed_size:
            raise ValueError("decompressed size {} vs {}".format(len(res), decompressed_size))
        return res



class CABAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "CAB"
    regexp = r"MSCF..\x00\x00....\x00\x00\x00\x00"

    def __init__(self):
        FileTypeAnalyzer.__init__(self)
        self.header = None
        self.folders = []
        self.data = []
        self.filesystem = {}

    def unpack(self, vfile, password=None):
        cffile = self.filesystem[vfile.path]
        folder = self.folders[cffile["FolderIndex"]]
        file_start = cffile["OffsetFolderStart"]
        folder_start = folder["DataStart"]
        file_size = cffile["UncompressedSize"]
        compression = folder["CompressionMethod"]
        data_count = folder["DataCount"]

        if compression == 0:
            decompressor = Decompressor()
        elif compression == 1:
            decompressor = ZlibDecompressor()
        elif compression == 3:
            decompressor = LzxDecompressor(folder["CompressionLevel"])
        else:
            raise ValueError("Unsupported compression method: 0x{:x}".format(compression))
        seen = None
        total_size = 0
        unpacked = []
        for cfdata in self.data:
            if cfdata.offset == folder_start:
                seen = 0
            buf = decompressor.decompress_one_block(cfdata["Data"], cfdata["UncompressedSize"])
            if seen is not None:
                seen += 1
                if total_size + len(buf) <= file_start:
                    file_start -= len(buf)
                else:
                    unpacked.append(buf)
                    total_size += len(buf)
                if seen >= data_count or total_size >= file_size + file_start:
                    break
        unpacked = b"".join(unpacked)
        return unpacked[file_start:file_start+file_size]

    def parse(self, hint):
        ch = yield CFHeader(category=Type.HEADER)
        reserved = 0
        start = self.tell()
        if "ReserveFolder" in ch:
            reserved = ch["ReserveFolder"]
        for i in range(ch["NumberOfFolders"]):
            f = yield CFFolder(reserved, name="CFFolder", category=Type.HEADER)
            self.folders.append(f)
        if not self.folders:
            raise FatalError("No folder")
        self.confirm()
        self.add_section("Folders", start, self.tell() - start)
        self.jump(ch["FirstFile"])
        start = self.tell()
        for i in range(ch["NumberOfFiles"]):
            folder_index, = struct.unpack("<H", self.read(self.tell() + 8, 2))
            if folder_index < len(self.folders):
                parent = self.folders[folder_index]
            elif folder_index == 0xffff or folder_index == 0xfffd:
                parent = self.folders[0]
            elif folder_index == 0xfffe:
                parent = self.folders[-1]
            else:
                parent = 0
            f = yield CFFile(parent=parent, category=Type.HEADER)
            self.filesystem[f["FileName"]] = f
            self.add_file(f["FileName"], f["UncompressedSize"], "unpack")
        self.add_section("Files", start, self.tell() - start)
        reserved = 0
        if "ReserveData" in ch:
            reserved = ch["ReserveData"]
        for i, folder in enumerate(self.folders):
            self.jump(folder["DataStart"])
            start = self.tell()
            for j in range(folder["DataCount"]):
                d = yield CFData(reserved, name="CFData", parent=folder, category=Type.DATA)
                self.data.append(d)
            self.add_section("Folder#{}".format(i), start, self.tell() - start)
        
