from filetypes.base import *
import malcat 
import io

class UnknownHeader(Struct):

    def parse(self):
        yield UInt16(name="Crc", comment="crc of data")
        size = yield UInt16(name="Size", comment="size of data")
        start = len(self)
        yield UInt8(name="Type", values=[
            ( "MAIN", 0),
            ( "FILE32", 1),
            ( "RECOVERY32", 2),
            ( "FILE64", 3),
            ( "RECOVERY64A", 4),
            ( "RECOVERY64B", 5),
            ])
        flags = yield BitsField(
            Bit(name="AddSize", comment="AddSize field present"),
            Bit(name="64Bits", comment="64bits size fields"),
            NullBits(14),
            name="Flags", comment="header flags")
        if flags["AddSize"]:
            if flags["64Bits"]:
                szfld = UInt64
            else:
                szfld = UInt32
            yield szfld(name="PackedSize", comment="Compressed data size")
        done = len(self) - start
        if size > done:
            yield Bytes(size-done, name="Data", comment="stream chunk")


class MainHeader(Struct):

    def parse(self):
        yield UInt16(name="Crc", comment="crc of data")
        size = yield UInt16(name="Size", comment="size of data")
        start = len(self)
        yield UInt8(name="Type", values=[
            ( "MAIN", 0),
            ( "FILE32", 1),
            ( "RECOVERY32", 2),
            ( "FILE64", 3),
            ( "RECOVERY64A", 4),
            ( "RECOVERY64B", 5),
            ])
        flags = yield BitsField(
            Bit(name="AddSize", comment="AddSize field present"),
            Bit(name="Comment", comment="Comment field present"),
            Bit(name="64Bits", comment="64bits size fields"),
            NullBits(5),
            Bit(name="V2Format", comment="ACE 2.0 format"),
            Bit(name="Sfx", comment="self-extracting archive"),
            Bit(name="LimitSfxJr", comment="dict size limited to 256K"),
            Bit(name="MultiVolume", comment="archive has multiple volumes"),
            Bit(name="Advert", comment="advert string present"),
            Bit(name="Recovery", comment="recovery record present"),
            Bit(name="Locked", comment="archive is locked"),
            Bit(name="Solid", comment="archive is solid"),
            name="Flags", comment="header flags")
        yield String(7, name="Signature", comment="magic")
        yield UInt8(name="EVersion")
        yield UInt8(name="CVersion")
        yield UInt8(name="Host", values=[
            ("HOST_MSDOS", 0),
            ("HOST_OS2", 1),
            ("HOST_WIN32", 2),
            ("HOST_UNIX", 3),
            ("HOST_MAC_OS", 4),
            ("HOST_WIN_NT", 5),
            ("HOST_PRIMOS", 6),
            ("HOST_APPLE_GS", 7),
            ("HOST_ATARI", 8),
            ("HOST_VAX_VMS", 9),
            ("HOST_AMIGA", 10),
            ("HOST_NEXT", 11),
            ("HOST_LINUX", 12),
            ])
        yield UInt8(name="Volume")
        yield Timestamp(name="CreationTime")
        yield Unused(8, name="Reserved")
        if flags["Advert"]:
            advsz = yield UInt8(name="AdvertSize", comment="size of advert string")
            s = yield String(advsz, zero_terminated=False, name="Advert", comment="advertising string")
        if flags["Comment"]:
            advsz = yield UInt8(name="CommentSize", comment="size of comment string")
            s = yield String(advsz, zero_terminated=False, name="Comment", comment="advertising string")
        done = len(self) - start
        if done < size:
            yield Bytes(size-done, name="Overlay", comment="overlay")    


class FileHeader(Struct):

    def parse(self):
        yield UInt16(name="Crc", comment="crc of data")
        size = yield UInt16(name="Size", comment="size of data")
        start = len(self)
        tp = yield UInt8(name="Type", values=[
            ( "MAIN", 0),
            ( "FILE32", 1),
            ( "RECOVERY32", 2),
            ( "FILE64", 3),
            ( "RECOVERY64A", 4),
            ( "RECOVERY64B", 5),
            ])
        is64 = tp == 3
        flags = yield BitsField(
            Bit(name="AddSize", comment="AddSize field present"),
            Bit(name="Comment", comment="Comment field present"),
            Bit(name="64Bits", comment="64bits size fields"),
            NullBits(7),
            Bit(name="NTSecurity", comment="NTFS security data present"),
            NullBits(1),
            Bit(name="ContPrev", comment="continued from previous volume"),
            Bit(name="Recovery", comment="recovery record present"),
            Bit(name="Password", comment="password encrypted"),
            Bit(name="Solid", comment="archive is solid"),
            name="Flags", comment="header flags")
        if flags["64Bits"] != is64:
            raise FatalError("incoherent 64Bits flag")
        if is64:
            szfld = UInt64
        else:
            szfld = UInt32
        yield szfld(name="PackedSize", comment="Compressed data size")
        yield szfld(name="OriginalSize", comment="Uncompressed file size")
        yield Timestamp(name="CreationTime")
        yield BitsField(
            Bit(name="ReadOnly"),
            Bit(name="Hidden"),
            Bit(name="System"),
            Bit(name="VolumeId"),
            Bit(name="Directory"),
            Bit(name="Archive"),
            Bit(name="Device"),
            Bit(name="Normal"),
            Bit(name="Temporary"),
            Bit(name="Sparse"),
            Bit(name="ReparsePoint"),
            Bit(name="Compressed"),
            Bit(name="Offline"),
            Bit(name="NotIndexed"),
            Bit(name="Encrypted"),
            Bit(name="IntegrityStream"),
            Bit(name="Virtual"),
            Bit(name="NoScrub"),
            Bit(name="EA"),
            NullBits(13),
            name="Attributes", comment="file characteristics")
        yield UInt32(name="Crc32", comment="file crc32") 
        yield UInt8(name="CompressionType", values=[
            ("COMP_STORED", 0),
            ("COMP_LZ77", 1),
            ("COMP_BLOCKED", 2),
            ])
        yield UInt8(name="CompressionQual")
        yield UInt16(name="Parameters")
        yield Unused(2, name="Reserved")
        fnsz = yield UInt16(name="FilenameSize")
        yield String(fnsz, zero_terminated=False, name="FileName")
        done = len(self) - start
        if done < size:
            yield Bytes(size-done, name="Overlay", comment="overlay")    



SIG_TO_CHUNK = {
        0: MainHeader,
        1: FileHeader,
        3: FileHeader,
}


class ACEAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "ACE"
    regexp = r"(?<=....\x00..)\*\*ACE\*\*"

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

    def open(self, vfile, password=None):
        import acefile
        filelike = io.BytesIO(self.read(0, self.size()))
        with acefile.open(filelike) as f:
            return f.read(vfile.path)

    def parse(self, hint):
        self.filesystem = {}
        hdr = yield MainHeader()
        self.confirm()

        if hdr["Flags"]["Advert"]:
            self.add_metadata("Advertising", hdr["Advert"])
        if hdr["Flags"]["Comment"]:
            self.add_metadata("Comment", hdr["Comment"])

        seen_end = False
        while not self.eof() and not seen_end:
            if self.remaining() < 7:
                break
            typ = self.read(self.tell() + 4, 1)[0]
            chunk_class = SIG_TO_CHUNK.get(typ, UnknownHeader)
            chk = yield chunk_class()
            chk_name = "Chunk"
            start = chk.offset
            if "FileName" in chk:
                chk_name = chk["FileName"]
                self.filesystem[chk_name] = chk
                self.add_file(chk_name, chk["OriginalSize"], "open")
            if "PackedSize" in chk:
                yield Bytes(chk["PackedSize"], name="Data", category=Type.DATA)
                size = self.tell() - start
                self.add_section(chk_name, start, size)
            

       
        

