import pefile
import sys
import binascii
import argparse

class icon_rule_maker():
    def __init__(self):
        self.pe = pefile.PE(args.path)
        self.EntryPoint = self.pe.OPTIONAL_HEADER.AddressOfEntryPoint
        self.ImageBase = self.pe.OPTIONAL_HEADER.ImageBase
        self.section_list = {}
        self.result = ""
        self.count = 1

        for section in self.pe.sections:
            self.section_list[section.Name.decode("utf-8").replace("\x00","")] = [hex(section.VirtualAddress), hex(section.SizeOfRawData), hex(section.PointerToRawData)]

    def make_icon(self, start, end, path):
        for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries:
            resource_type = entry.name
            if resource_type is None:
                resource_type = pefile.RESOURCE_TYPE.get(entry.struct.Id)

            for directory in entry.directory.entries:
                for resource in directory.directory.entries:
                    name = str(resource_type)
                    if name in "RT_ICON":
                        name = str(resource_type)
                        offset = resource.data.struct.OffsetToData
                        size = resource.data.struct.Size
                        RVA_ = int(self.section_list['.rsrc'][0],16) - int(self.section_list['.rsrc'][2],16)
                        print(name, hex(offset), hex(size))

                        real_offset = hex(offset - RVA_)
                        print(hex(offset), real_offset)

                        f = open(args.path, "rb")
                        f.seek(int(real_offset,16))

                        data = binascii.hexlify(f.read(size))[start:end].decode("utf-8")
                        f.close()

                        count = 0
                        for i in data:
                            if i == "0":
                                count += 1

                        print(data, count)

                        if not count == 600:
                            self.result += "rule icon_" + str(self.count) + "{ strings: \n $a = {"
                            self.count +=1
                            for i in range(0, len(data), 2):
                                self.result += str(data[i]) + str(data[i+1]) + " "
                            self.result += "}\n condition: \n all of them \n }"

        f = open("rule.yar","w")
        f.write(self.result)
        f.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-s","--start", help="start icon offset")
    parser.add_argument("-e","--end", help="end icon offset")
    parser.add_argument("-t","--type", help="hex or int")
    parser.add_argument("-p","--path", help="Binary Path")
    args = parser.parse_args()
    a = icon_rule_maker()

    if not args.start and not args.end and not args.type:
        a.make_icon(-600, None, args.path)
    else:
        if args.type == "hex":
            a.make_icon(int(args.start, 16), int(args.end, 16),args.path)
        elif args.type == "int":
            a.make_icon(int(args.start), int(args.end),args.path)
        else:
            sys.exit()