from filetypes.base import *
import malcat
import struct

# https://codedread.github.io/bitjs/docs/unrar.html

class Block(Struct):

    def parse(self):
        yield UInt16(name="HeadCrc", comment="block crc16")
        headtype = yield UInt8(name="HeadType", comment="block type")
        flags = yield UInt16(name="HeadFlags")
        size = yield UInt16(name="HeadSize", comment="block flags")
        if flags & 0x8000:
            add_size = yield UInt32(name="AdditionalSize", comment="AdditionalSize")
            size += add_size
        if size < len(self):
            raise FatalError("size too small: {}".format(size))
        if size > len(self):
            yield Bytes(size - len(self), name="Data")


class MainBlock(Struct):

    def parse(self):
        yield UInt16(name="HeadCrc", comment="block crc16")
        headtype = yield UInt8(name="HeadType", comment="block type")
        flags = yield BitsField(
            Bit(name="MHD_VOLUME"),
            Bit(name="MHD_COMMENT"),
            Bit(name="MHD_LOCK"),
            Bit(name="MHD_SOLID"),
            Bit(name="MHD_PACK_COMMENT"),
            Bit(name="MHD_AV"),
            Bit(name="MHD_PROTECT"),
            Bit(name="MHD_PASSWORD"),
            Bit(name="MHD_PACK_COMMENT"),
            Bit(name="MHD_FIRSTVOLUME"),
            Bit(name="MHD_ENCRYPTVER"),
            NullBits(4),
            Bit(name="MHD_ADDITIONAL_SIZE"),
            name="Flags")
        size = yield UInt16(name="HeadSize", comment="block flags")
        if flags["MHD_ADDITIONAL_SIZE"]:
            add_size = yield UInt32(name="AdditionalSize", comment="AdditionalSize")
            size += add_size
        if size < len(self):
            raise FatalError("size too small: {}".format(size))
        if size > len(self):
            yield Bytes(size - len(self), name="Data")

class FileBlock(Struct):

    def parse(self):
        yield UInt16(name="HeadCrc", comment="block crc16")
        headtype = yield UInt8(name="HeadType", comment="block type")
        flags = yield BitsField(
            Bit(name="LHD_SPLIT_BEFORE"),
            Bit(name="LHD_SPLIT_AFTER"),
            Bit(name="LHD_PASSWORD"),
            Bit(name="LHD_COMMENT"),
            Bit(name="LHD_SOLID"),
            NullBits(3),
            Bit(name="LHD_LARGE", comment="64 bits are used to describe the packed and unpacked size"),
            Bit(name="LHD_UNICODE", comment="filename is Unicode"),
            Bit(name="LHD_SALT", comment="64-bit salt value is present"),
            Bit(name="LHD_VERSION"),
            Bit(name="LHD_EXTTIME", comment="ExtTime Structure is present in the FILE_HEAD header"),
            Bit(name="LHD_EXTFLAGS"),
            NullBits(1),
            Bit(name="LHD_ADDITIONAL_SIZE"),
            name="HeadFlags")
        size = yield UInt16(name="HeadSize", comment="block flags")
        packsize = yield UInt32(name="PackSize", comment="")
        if flags["LHD_ADDITIONAL_SIZE"]:
            size += packsize
        yield UInt32(name="UnpSize", comment="")
        yield UInt8(name="HostOS", comment="")
        yield UInt32(name="FileCRC", comment="")
        yield Timestamp(name="FileTime", comment="")
        yield UInt8(name="UnpVer", comment="")
        yield UInt8(name="Method", comment="")
        namesz = yield UInt16(name="NameSize", comment="")
        yield UInt32(name="FileAttr", comment="")
        if flags["LHD_LARGE"]:
            yield UInt32(name="HighPackSize", comment="")
            yield UInt32(name="HighUnpSize", comment="")
        if flags["LHD_UNICODE"]:
            yield StringUtf16le(namesz//2, zero_terminated=False, name="FileName")
        else:
            yield StringUtf8(namesz, zero_terminated=False, name="FileName")
        if flags["LHD_SALT"]:
            yield Bytes(8, name="Salt")
        if size < len(self):
            raise FatalError("size too small: {}".format(size))
        if size > len(self):
            yield Bytes(size - len(self), name="PackedData")


BLOCK_TYPES = {
        0x72: ("Mark", Block),
        0x73: ("Main", MainBlock),
        0x74: ("File", FileBlock),
        0x75: ("Comm", Block),
        0x76: ("Av", Block),
        0x77: ("Sub", Block),
        0x78: ("Protect", Block),
        0x79: ("Sign", Block),
        0x7a: ("NewSub", FileBlock),
        0x7b: ("End", Block),
}

class RarAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "RAR"
    regexp = r"Rar!\x1a\x07\x00"


    def parse(self, hint):
        mb = yield Block(name="Marker", category=Type.HEADER)
        if mb["HeadType"] != 0x72:
            raise FatalError("Not a valid marker block")
        while mb["HeadType"] != 0x7b and self.remaining():
            if mb["HeadType"] == 0x74:
                self.add_section(mb["FileName"], mb.offset, mb.size)
                self.add_file(mb["FileName"], mb.size)
            type_byte, = struct.unpack("B", self.read(self.tell() + 2, 1))
            if type_byte not in BLOCK_TYPES:
                raise FatalError("Invalid block type 0x{:x}".format(type_byte))
            block_name, block_class = BLOCK_TYPES[type_byte]
            mb = yield block_class(name=block_name, category=Type.DATA)


#####################################################################

TYPE2NAME5 = {
        0x72: "Mark",
        0x73: "Main",
        0x74: "File",
        0x75: "Comm",
        0x76: "Av",
        0x77: "Sub",
        0x78: "Protect",
        0x79: "Sign",
        0x7a: "NewSub",
        0x7b: "End",
}




class Block5(Struct):

    def parse(self):
        yield UInt32(name="HeadCrc", comment="CRC32 of header data starting from Header size field and up to and including the optional extra area.")
        size = yield VarUInt64(name="HeadSize", comment="size of header data starting from Header type field and up to and including the optional extra area. This field must not be longer than 3 bytes in current implementation, resulting in 2 MB maximum header size")
        off = len(self)
        headtype = yield VarUInt64(name="HeadType", comment="block type")
        flags = yield VarUInt64(name="HeadFlags", comment="block flags")
        if flags > 0x80:
            raise FatalError("Invalid flag value: {:x}".format(flags))
        if flags & 0x0001:
            extra_size = yield VarUInt64(name="ExtraAreaSize", comment="size of extra area. Optional field, present only if 0x0001 header flag is set")
        else:
            extra_size = 0
        if flags & 0x0002:
            data_size = yield VarUInt64(name="DataSize", comment="size of data area. Optional field, present only if 0x0002 header flag is set")
        else:
            data_size = 0
        done = len(self) - off
        content = size - (done + extra_size)
        if content < 0:
            raise FatalError("Invalid block sizes: header_size={} extra_size={} data_size={}".format(size, extra_size, data_size))
        yield Bytes(content, name="OtherFields", comment="fields specific for current block type")
        if extra_size:
            yield Bytes(extra_size, name="ExtraArea", comment="optional area containing additional header fields")
        if data_size:
            yield Bytes(data_size, name="DataArea", comment="optional data area")            

class EncryptionHeader(Struct):

    def parse(self):
        yield UInt32(name="HeadCrc", comment="CRC32 of header data starting from Header size field and up to and including the optional extra area.")
        size = yield VarUInt64(name="HeadSize", comment="size of header data starting from Header type field and up to and including the optional extra area. This field must not be longer than 3 bytes in current implementation, resulting in 2 MB maximum header size")
        off = len(self)
        headtype = yield VarUInt64(name="HeadType", comment="block type")
        flags = yield VarUInt64(name="HeadFlags", comment="block flags")
        if flags > 0x80:
            raise FatalError("Invalid flag value: {:x}".format(flags))
        if flags & 0x0001:
            raise FatalError("Invalid flag value: {:x}".format(flags))
        if flags & 0x0002:
            raise FatalError("Invalid flag value: {:x}".format(flags))
        yield VarUInt64(name="EnryptionVersion", comment="encryption algorithm")
        encryption_flags = yield VarUInt64(name="EnryptionFlags", comment="encryption flags")
        yield UInt8(name="KDFCount", comment="Binary logarithm of iteration number for PBKDF2 function")
        yield Bytes(16, name="salt", comment="salt value used globally for all encrypted archive headers")
        if encryption_flags & 0x01:
            yield Bytes(12, name="CheckValue", comment="value used to verify the password validity")
        done = len(self) - off
        content = size - done
        if content != 0:
            raise FatalError("Invalid block sizes: header_size={} extra_size={} data_size={}".format(size, extra_size, data_size))
        yield Bytes(content, name="OtherFields", comment="fields specific for current block type")

class FileOrServiceHeader(Struct):

    def parse(self):
        yield UInt32(name="HeadCrc", comment="CRC32 of header data starting from Header size field and up to and including the optional extra area.")
        size = yield VarUInt64(name="HeadSize", comment="size of header data starting from Header type field and up to and including the optional extra area. This field must not be longer than 3 bytes in current implementation, resulting in 2 MB maximum header size")
        off = len(self)
        headtype = yield VarUInt64(name="HeadType", comment="block type")
        flags = yield VarUInt64(name="HeadFlags", comment="block flags")
        if flags > 0x80:
            raise FatalError("Invalid flag value: {:x}".format(flags))
        if flags & 0x0001:
            extra_size = yield VarUInt64(name="ExtraAreaSize", comment="size of extra area. Optional field, present only if 0x0001 header flag is set")
        else:
            extra_size = 0
        if flags & 0x0002:
            data_size = yield VarUInt64(name="DataSize", comment="size of data area. Optional field, present only if 0x0002 header flag is set")
        else:
            data_size = 0
        file_flags = yield VarUInt64(name="FileFlags", comment="file-related flags")
        unpacked_size = yield VarUInt64(name="UnpackedSize", comment="unpacked file or service data size")
        attributes = yield VarUInt64(name="Attributes", comment="operating system specific file attributes in case of file header, might be either used for data specific needs or just reserved and set to 0 for service header")
        if file_flags & 2:
            yield UInt32(name="Mtime", comment="file modification time in Unix time format. Optional, present if 0x0002 file flag is set")
        if file_flags & 4:
            yield UInt32(name="DataCrc", comment="CRC32 of unpacked file or service data. For files split between volumes it contains CRC32 of file packed data contained in current volume for all file parts except the last. Optional, present if 0x0004 file flag is set")
        yield VarUInt64(name="CompressionInfo", comment="Lower 6 bits (0x003f mask) contain the version of compression algorithm, resulting in possible 0 - 63 values. Current version is 0. 7th bit (0x0040) defines the solid flag. If it is set, RAR continues to use the compression dictionary left after processing preceding files. It can be set only for file headers and is never set for service headers. Bits 8 - 10 (0x0380 mask) define the compression method. Currently only values 0 - 5 are used. 0 means no compression. Bits 11 - 14 (0x3c00) define the minimum size of dictionary size required to extract data. Value 0 means 128 KB, 1 - 256 KB, ..., 14 - 2048 MB, 15 - 4096 MB.")
        yield UInt8(name="HostOS", comment="operating system used to create the archive")
        name_length = yield VarUInt64(name="NameLength", comment="header name length")
        yield StringUtf8(name_length, name="Name", comment="For file header this is a name of archived file. Forward slash character is used as the path separator both for Unix and Windows names. Backslashes are treated as a part of name for Unix names and as invalid character for Windows file names. Type of name is defined by Host OS field. For service header this field contains a name of service header")
        done = len(self) - off
        content = size - (done + extra_size)
        if content != 0:
            raise FatalError("Invalid block sizes: header_size={} extra_size={} data_size={}".format(size, extra_size, data_size))
        if extra_size:
            yield Bytes(extra_size, name="ExtraArea", comment="optional area containing additional header fields")
        if data_size:
            yield Bytes(min(data_size, self.remaining()), name="DataArea", comment="optional data area")                      


TYPEINFOS = {
        1: ("MainHeader", Block5, Type.HEADER),
        2: ("FileHeader", FileOrServiceHeader, Type.DATA),
        3: ("ServiceHeader", FileOrServiceHeader, Type.META),
        4: ("EncryptionHeader", EncryptionHeader, Type.HEADER),
        5: ("EOFHeader", Block5, Type.HEADER),
}


class Rar5Analyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "RAR5"
    regexp = r"Rar!\x1a\x07\x01\x00"

 
    def parse(self, hint):
        yield Bytes(8, name="Signature", category=Type.HEADER)
        first = True
        while self.remaining():
            off = self.tell()
            head_size, head_size_enc = VarUInt64.parse(self, off + 4)
            head_type, _ = VarUInt64.parse(self, off + 4 + head_size_enc)
            head_name, cls, category = TYPEINFOS.get(head_type, (None, None, None))
            if not head_name:
                raise FatalError("Invalid block type: {}".format(head_type))
            el = yield cls(name=head_name, category=category)
            if head_type == 2:
                self.add_section("Name" in el and el["Name"] or "???", el.offset, el.size)
                self.add_file("Name" in el and el["Name"] or "???", el.size)
            if head_type == 5 or head_type == 4:
                break
            if first and head_type != 4 and head_type != 1:
                raise FatalError("Invalid head tye {}".format(head_type))
            if first:
                # we got at least one valid block, so this is a valid rar5 file
                self.confirm()
            first = False
            if head_type == 4:
                break   # we can't read anything after the encryption header
            

       
