from filetypes.base import *
import malcat
from filetypes.BMP import BMPRawAnalyzer
from filetypes.PNG import PNGAnalyzer
import struct


class IconDirectory(Struct):

    def parse(self):
        yield UInt16(name="Signature", comment="signature (always 0)")
        yield UInt16(name="Type", comment="type (should be 1)")
        yield UInt16(name="NumberOfItems", comment="number of items")


class IconDirectoryEntry(Struct):

    def parse(self):
        yield UInt8(name="Width", comment="width")
        yield UInt8(name="Height", comment="height")
        yield UInt8(name="NumberOfColors", comment="number of colors")
        yield UInt8(name="Reserved", comment="reserved")
        yield UInt16(name="Planes", comment="color planes (=1)")
        yield UInt16(name="Bpp", comment="bits per pixel")
        yield UInt32(name="Size", comment="content size in bytes")
        yield Offset32(name="Offset", comment="data offset")




class ICOAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.IMAGE
    name = "ICO"
    regexp = r"(?<=\x00\x00)\x01\x00[\x01-\x14]\x00(?:\x00\x00|\x10\x10|\x20\x20|\x30\x30|\x40\x40|\x48\x48|\x60\x60|\x80\x80)[\x00\x10]\x00[\x00\x01\x04][\x00\x08\x18\x20]"

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

    def __init__(self):
        FileTypeAnalyzer.__init__(self)

    
    def open(self, vfile, password=None):
        index = int(vfile.path.split("_")[1])
        return self.read(self["Items"][index]["Offset"], self["Items"][index]["Size"])
    

    def parse(self, hint):
        idir = yield IconDirectory(category=Type.HEADER)
        nb = idir["NumberOfItems"]
        items = yield Array(nb, IconDirectoryEntry(), name="Items", category=Type.HEADER)
        for i in range(nb):
            off = items[i]["Offset"]
            sz = items[i]["Size"]
            self.add_file("item_{}".format(i), sz, "open")
            self.jump(off)
            if sz > self.remaining():
                raise FatalError
            self.add_section("Item{:d}".format(i), off, sz)

            self.jump(off)
            if sz > 3 and self.read(off, 4) == b"\x89PNG":
                png = PNGAnalyzer()
                for struct, _ in self.subparse(png, sz, hint="ICO"):
                    struct.name = "Item.{:d}.PNG.{}".format(i, struct.name)
                    struct.parent = items
                    el = yield struct
                if png.remaining():
                    el = yield Unused(png.remaining(), name="Item.{:d}.Overlay".format(i), category=Type.ANOMALY, parent=items)
            else:
                bmp = BMPRawAnalyzer()
                for struct, _ in self.subparse(bmp, sz, hint="ICO"):
                    struct.name = "Item.{:d}.BMP.{}".format(i, struct.name)
                    struct.parent = items
                    el = yield struct
                if not bmp.eof():
                    el = yield Unused(bmp.remaining(), name="Item.{:d}.Overlay".format(i), category=Type.ANOMALY, parent=items)
            
